# -*- coding: utf-8 -*-
#
# Part of ChipBench (AGPLv3). SPDX-FileCopyrightText: 2026 Starting Line Productions LLC
# Analyze8051.py -- force-disassemble a raw 8051 binary and re-run analysis.
# Raw binary imports leave no entry points, so Ghidra finds 0 functions.
# This script disassembles from the reset + interrupt vectors, re-analyzes,
# then sweeps remaining undefined bytes in populated code regions.

from ghidra.app.cmd.disassemble import DisassembleCommand
from ghidra.program.model.address import AddressSet
from ghidra.app.plugin.core.analysis import AutoAnalysisManager

prog = currentProgram
codeSpace = prog.getAddressFactory().getAddressSpace("CODE")
if codeSpace is None:
    codeSpace = prog.getAddressFactory().getDefaultAddressSpace()

def A(x):
    return codeSpace.getAddress(x)

# 8051 interrupt vectors (standard set; enhanced variants may extend)
vectors = [0x0000, 0x0003, 0x000B, 0x0013, 0x001B, 0x0023, 0x002B, 0x0033, 0x003B]

for v in vectors:
    cmd = DisassembleCommand(A(v), None, True)
    cmd.applyTo(prog, monitor)

instr = prog.getListing().getNumInstructions()
print("After vector disassembly:", instr, "instructions")

mgr = AutoAnalysisManager.getAnalysisManager(prog)
mgr.reAnalyzeAll(None)
mgr.startAnalysis(monitor)

instr = prog.getListing().getNumInstructions()
funcs = prog.getFunctionManager().getFunctionCount()
print("After analysis:", instr, "instructions,", funcs, "functions")

# Sweep remaining undefined bytes in all populated regions, 1KB at a time.
# String islands (0x9250-0x9F00, 0xEF00-0xF400) stay data; the rest is code.
listing = prog.getListing()
memory = prog.getMemory()
SKIP = set(range(0x9200, 0x9F00, 0x400)) | set(range(0xEE00, 0x10000, 0x400))
for base in range(0x0000, 0xF400, 0x400):
    if base in SKIP:
        continue
    blk = memory.getBlock(A(base))
    if blk is None:
        continue
    # skip fully-erased ranges
    erased = True
    for i in range(base, base + 0x400, 32):
        b = memory.getByte(A(i))
        if (b & 0xFF) != 0xFF:
            erased = False
            break
    if erased:
        continue
    rng = AddressSet(A(base), A(min(base + 0x3FF, 0xF3FF)))
    sweep = DisassembleCommand(A(base), rng, True)
    sweep.applyTo(prog, monitor)

mgr.reAnalyzeAll(None)
mgr.startAnalysis(monitor)

instr = prog.getListing().getNumInstructions()
funcs = prog.getFunctionManager().getFunctionCount()
print("Final:", instr, "instructions,", funcs, "functions")