unstrip 1.1.0

use std::io::{self, Write};

use crate::pclntab::Function;
use crate::types::{KindData, KindName, StructField, Type};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Target {
    Ida,
    Ghidra,
    BinaryNinja,
}

/// Emit a script that, run inside the target reverse engineering tool,
/// applies every recovered function name (with file:line as a function
/// comment) and every recovered struct type (as a C declaration the tool
/// can parse).
///
/// Pass `types` as an empty slice to emit functions only. Pass
/// `signatures = None` to emit bare function names without the recovered
/// Go-syntax signature in the per-function comment.
pub fn write_script<W: Write>(
    w: &mut W,
    target: Target,
    funcs: &[Function],
    types: &[Type],
    signatures: Option<&std::collections::HashMap<u64, String>>,
    generator: &str,
) -> io::Result<()> {
    match target {
        Target::Ida => write_ida(w, funcs, types, signatures, generator),
        Target::Ghidra => write_ghidra(w, funcs, types, signatures, generator),
        Target::BinaryNinja => write_binja(w, funcs, types, signatures, generator),
    }
}

fn sig_for(
    f: &Function,
    signatures: Option<&std::collections::HashMap<u64, String>>,
) -> Option<String> {
    signatures.and_then(|m| m.get(&f.address)).cloned()
}

fn write_ida<W: Write>(
    w: &mut W,
    funcs: &[Function],
    types: &[Type],
    signatures: Option<&std::collections::HashMap<u64, String>>,
    generator: &str,
) -> io::Result<()> {
    writeln!(
        w,
        "# Generated by {generator}. Run inside IDA: File -> Script File... (Python3)"
    )?;
    writeln!(w, "import ida_funcs, ida_name, ida_bytes, ida_typeinf, idc")?;
    writeln!(w)?;
    emit_stats_preamble(w)?;
    writeln!(w, "def _func(addr, name, comment):")?;
    writeln!(w, "    try:")?;
    writeln!(w, "        ok = ida_funcs.add_func(addr)")?;
    writeln!(
        w,
        "        ok = ida_name.set_name(addr, name, ida_name.SN_FORCE | ida_name.SN_NOCHECK) and ok"
    )?;
    writeln!(w, "        if comment:")?;
    writeln!(w, "            idc.set_func_cmt(addr, comment, 0)")?;
    writeln!(w, "        if ok:")?;
    writeln!(w, "            _stats['func_ok'] += 1")?;
    writeln!(w, "        else:")?;
    writeln!(w, "            _stats['func_fail'] += 1")?;
    writeln!(
        w,
        "            _stats['failures'].append(('func', '0x{{:x}} {{}}'.format(addr, name), 'IDA refused add_func or set_name'))"
    )?;
    writeln!(w, "    except Exception as e:")?;
    writeln!(w, "        _stats['func_fail'] += 1")?;
    writeln!(
        w,
        "        _stats['failures'].append(('func', '0x{{:x}} {{}}'.format(addr, name), repr(e)))"
    )?;
    writeln!(w)?;
    writeln!(w, "def _struct(decl):")?;
    writeln!(w, "    try:")?;
    writeln!(w, "        til = ida_typeinf.get_idati()")?;
    writeln!(
        w,
        "        rc = ida_typeinf.parse_decls(til, decl, None, ida_typeinf.PT_TYP)"
    )?;
    // ida_typeinf.parse_decls returns the number of errors, so 0 means
    // success. Any positive value indicates the parser rejected at least
    // one struct in the supplied decl block.
    writeln!(w, "        if rc == 0:")?;
    writeln!(w, "            _stats['struct_ok'] += 1")?;
    writeln!(w, "        else:")?;
    writeln!(w, "            _stats['struct_fail'] += 1")?;
    writeln!(
        w,
        "            _stats['failures'].append(('struct', decl.split('{{')[0].strip()[:80], 'IDA parse_decls returned {{}} errors'.format(rc)))"
    )?;
    writeln!(w, "    except Exception as e:")?;
    writeln!(w, "        _stats['struct_fail'] += 1")?;
    writeln!(
        w,
        "        _stats['failures'].append(('struct', decl.split('{{')[0].strip()[:80], repr(e)))"
    )?;
    writeln!(w)?;

    for f in funcs {
        let name = sanitize_label(&f.name);
        let sig = sig_for(f, signatures);
        let comment = function_comment(f, sig.as_deref());
        writeln!(
            w,
            "_func(0x{:x}, {}, {})",
            f.address,
            escape_py(&name),
            escape_py(&comment),
        )?;
    }

    if !types.is_empty() {
        writeln!(w)?;
        writeln!(
            w,
            "# Recovered struct types ({} total)",
            count_struct_types(types)
        )?;
        for decl in struct_c_decls(types) {
            writeln!(w, "_struct({})", escape_py(&decl))?;
        }
    }

    emit_stats_postamble(w, generator)?;
    Ok(())
}

fn write_ghidra<W: Write>(
    w: &mut W,
    funcs: &[Function],
    types: &[Type],
    signatures: Option<&std::collections::HashMap<u64, String>>,
    generator: &str,
) -> io::Result<()> {
    writeln!(
        w,
        "# Generated by {generator}. Run inside Ghidra: Window -> Script Manager (Python)."
    )?;
    writeln!(w, "# @author Riven Labs")?;
    writeln!(w, "# @category {generator}")?;
    writeln!(w, "# @keybinding")?;
    writeln!(w, "# @menupath Tools.{generator}.Apply recovered symbols")?;
    writeln!(w, "# @toolbar")?;
    writeln!(w, "# @runtime PyGhidra")?;
    writeln!(
        w,
        "# @description Apply {generator}-recovered Go function names, file:line comments, and struct types to the current program."
    )?;
    writeln!(w)?;
    writeln!(
        w,
        "from ghidra.program.model.symbol.SourceType import USER_DEFINED"
    )?;
    writeln!(w, "from ghidra.app.util.cparser.C import CParser")?;
    writeln!(
        w,
        "from ghidra.program.model.data import DataTypeConflictHandler"
    )?;
    writeln!(w)?;
    writeln!(w, "fm = currentProgram.getFunctionManager()")?;
    writeln!(w, "af = currentProgram.getAddressFactory()")?;
    writeln!(w, "dtm = currentProgram.getDataTypeManager()")?;
    writeln!(w, "listing = currentProgram.getListing()")?;
    writeln!(w)?;
    emit_stats_preamble(w)?;
    writeln!(w, "def _func(addr, name, comment):")?;
    writeln!(w, "    label = '0x{{:x}} {{}}'.format(addr, name)")?;
    writeln!(w, "    try:")?;
    writeln!(
        w,
        "        a = af.getDefaultAddressSpace().getAddress(addr)"
    )?;
    writeln!(w, "        f = fm.getFunctionAt(a)")?;
    writeln!(w, "        if f is None:")?;
    writeln!(
        w,
        "            f = fm.createFunction(name, a, None, USER_DEFINED)"
    )?;
    writeln!(w, "        else:")?;
    writeln!(w, "            f.setName(name, USER_DEFINED)")?;
    writeln!(w, "        if comment:")?;
    writeln!(w, "            f.setComment(comment)")?;
    writeln!(w, "        _stats['func_ok'] += 1")?;
    writeln!(w, "    except Exception as e:")?;
    writeln!(w, "        _stats['func_fail'] += 1")?;
    writeln!(
        w,
        "        _stats['failures'].append(('func', label, repr(e)))"
    )?;
    writeln!(w)?;
    writeln!(w, "def _struct(decl):")?;
    writeln!(w, "    label = decl.split('{{')[0].strip()[:80]")?;
    writeln!(w, "    try:")?;
    writeln!(w, "        parser = CParser(dtm)")?;
    writeln!(w, "        parsed = parser.parse(decl)")?;
    writeln!(w, "        if parsed is None:")?;
    writeln!(w, "            _stats['struct_fail'] += 1")?;
    writeln!(
        w,
        "            _stats['failures'].append(('struct', label, 'CParser returned None'))"
    )?;
    writeln!(w, "            return")?;
    writeln!(
        w,
        "        dtm.addDataType(parsed, DataTypeConflictHandler.REPLACE_HANDLER)"
    )?;
    writeln!(w, "        _stats['struct_ok'] += 1")?;
    writeln!(w, "    except Exception as e:")?;
    writeln!(w, "        _stats['struct_fail'] += 1")?;
    writeln!(
        w,
        "        _stats['failures'].append(('struct', label, repr(e)))"
    )?;
    writeln!(w)?;

    for f in funcs {
        let name = sanitize_label(&f.name);
        let sig = sig_for(f, signatures);
        let comment = function_comment(f, sig.as_deref());
        writeln!(
            w,
            "_func(0x{:x}, {}, {})",
            f.address,
            escape_py(&name),
            escape_py(&comment),
        )?;
    }

    if !types.is_empty() {
        writeln!(w)?;
        writeln!(
            w,
            "# Recovered struct types ({} total)",
            count_struct_types(types)
        )?;
        for decl in struct_c_decls(types) {
            writeln!(w, "_struct({})", escape_py(&decl))?;
        }
    }

    emit_stats_postamble(w, generator)?;
    Ok(())
}

fn write_binja<W: Write>(
    w: &mut W,
    funcs: &[Function],
    types: &[Type],
    signatures: Option<&std::collections::HashMap<u64, String>>,
    generator: &str,
) -> io::Result<()> {
    writeln!(w, "# Generated by {generator}. Run inside Binary Ninja:")?;
    writeln!(
        w,
        "#   bv = binaryninja.load('binary'); exec(open('this.py').read())"
    )?;
    writeln!(
        w,
        "from binaryninja import Symbol, SymbolType, Type as BNType"
    )?;
    writeln!(w)?;
    emit_stats_preamble(w)?;
    writeln!(w, "def _func(addr, name, comment):")?;
    writeln!(w, "    label = '0x{{:x}} {{}}'.format(addr, name)")?;
    writeln!(w, "    try:")?;
    writeln!(w, "        bv.add_function(addr)")?;
    writeln!(
        w,
        "        bv.define_user_symbol(Symbol(SymbolType.FunctionSymbol, addr, name))"
    )?;
    writeln!(w, "        if comment:")?;
    writeln!(w, "            f = bv.get_function_at(addr)")?;
    writeln!(w, "            if f is not None:")?;
    writeln!(w, "                f.comment = comment")?;
    writeln!(w, "        _stats['func_ok'] += 1")?;
    writeln!(w, "    except Exception as e:")?;
    writeln!(w, "        _stats['func_fail'] += 1")?;
    writeln!(
        w,
        "        _stats['failures'].append(('func', label, repr(e)))"
    )?;
    writeln!(w)?;
    writeln!(w, "def _struct(decl):")?;
    writeln!(w, "    label = decl.split('{{')[0].strip()[:80]")?;
    writeln!(w, "    try:")?;
    writeln!(w, "        types = bv.parse_types_from_string(decl)")?;
    writeln!(w, "        for name, typ in types.types.items():")?;
    writeln!(w, "            bv.define_user_type(name, typ)")?;
    writeln!(w, "        _stats['struct_ok'] += 1")?;
    writeln!(w, "    except Exception as e:")?;
    writeln!(w, "        _stats['struct_fail'] += 1")?;
    writeln!(
        w,
        "        _stats['failures'].append(('struct', label, repr(e)))"
    )?;
    writeln!(w)?;

    for f in funcs {
        let name = sanitize_label(&f.name);
        let sig = sig_for(f, signatures);
        let comment = function_comment(f, sig.as_deref());
        writeln!(
            w,
            "_func(0x{:x}, {}, {})",
            f.address,
            escape_py(&name),
            escape_py(&comment),
        )?;
    }

    if !types.is_empty() {
        writeln!(w)?;
        writeln!(
            w,
            "# Recovered struct types ({} total)",
            count_struct_types(types)
        )?;
        for decl in struct_c_decls(types) {
            writeln!(w, "_struct({})", escape_py(&decl))?;
        }
    }

    emit_stats_postamble(w, generator)?;
    Ok(())
}

/// Emit the shared Python preamble that initializes the per-script
/// success/failure counters. Every writer calls this once before the
/// per-symbol _func() and _struct() calls so the helpers have a place
/// to record what landed and what did not.
fn emit_stats_preamble<W: Write>(w: &mut W) -> io::Result<()> {
    writeln!(w, "_stats = {{")?;
    writeln!(w, "    'func_ok': 0,")?;
    writeln!(w, "    'func_fail': 0,")?;
    writeln!(w, "    'struct_ok': 0,")?;
    writeln!(w, "    'struct_fail': 0,")?;
    writeln!(w, "    'failures': [],")?;
    writeln!(w, "}}")?;
    writeln!(w)?;
    Ok(())
}

/// Emit the shared Python postamble that prints honest counts and a
/// per-failure breakdown. Called once at the bottom of every writer
/// after the last _func() / _struct() call.
fn emit_stats_postamble<W: Write>(w: &mut W, generator: &str) -> io::Result<()> {
    writeln!(w)?;
    writeln!(w, "_func_total = _stats['func_ok'] + _stats['func_fail']")?;
    writeln!(
        w,
        "_struct_total = _stats['struct_ok'] + _stats['struct_fail']"
    )?;
    writeln!(
        w,
        "print('{generator}: applied {{}}/{{}} symbols ({{}} failed)'.format(_stats['func_ok'], _func_total, _stats['func_fail']))"
    )?;
    writeln!(
        w,
        "print('{generator}: applied {{}}/{{}} struct types ({{}} failed)'.format(_stats['struct_ok'], _struct_total, _stats['struct_fail']))"
    )?;
    // Limit the per-failure dump to keep the Script Manager output
    // readable; surface enough to start debugging without dumping
    // thousands of lines on a bad export.
    writeln!(w, "if _stats['failures']:")?;
    writeln!(
        w,
        "    print('{generator}: first {{}} failures:'.format(min(20, len(_stats['failures']))))"
    )?;
    writeln!(w, "    for kind, key, reason in _stats['failures'][:20]:")?;
    writeln!(
        w,
        "        print('  {{}}  {{}}  {{}}'.format(kind, key, reason))"
    )?;
    writeln!(w, "    if len(_stats['failures']) > 20:")?;
    writeln!(
        w,
        "        print('  ... and {{}} more (see _stats[\"failures\"] for the full list)'.format(len(_stats['failures']) - 20))"
    )?;
    Ok(())
}

/// Sanitize a Go symbol name into a label the disassembler will accept.
/// Go names contain `.`, `(`, `)`, `*`, `[`, `]`, `/`, `,`, `<-`, etc.
/// IDA accepts most of these; Ghidra is stricter. We map the universally
/// rejected characters to `_` and keep `.` since both tools tolerate it.
fn sanitize_label(name: &str) -> String {
    let mut out = String::with_capacity(name.len());
    for ch in name.chars() {
        match ch {
            'a'..='z' | 'A'..='Z' | '0'..='9' | '_' | '.' | '$' => out.push(ch),
            // Replace structural chars with underscore. The original full name
            // ends up in the function comment so nothing is lost.
            '(' | ')' | '[' | ']' | '*' | '/' | ',' | ' ' | '-' | '<' | '>' | '{' | '}' | ':'
            | ';' | '"' | '\'' | '?' | '!' | '#' | '@' | '\\' | '+' | '=' | '&' | '|' | '^'
            | '~' | '%' => out.push('_'),
            _ => out.push('_'),
        }
    }
    // Disassemblers reject labels starting with a digit.
    if out
        .chars()
        .next()
        .map(|c| c.is_ascii_digit())
        .unwrap_or(false)
    {
        out.insert(0, '_');
    }
    if out.is_empty() {
        out.push('_');
    }
    out
}

/// Per-function comment: "go: <full name>\n<file>:<line>".
fn function_comment(f: &Function, signature: Option<&str>) -> String {
    let mut s = format!("go: {}", f.name);
    if let Some(sig) = signature {
        s.push_str(sig);
    }
    if let Some(file) = &f.file {
        s.push('\n');
        s.push_str(file);
        if let Some(line) = f.start_line {
            s.push(':');
            s.push_str(&line.to_string());
        }
    }
    s
}

fn count_struct_types(types: &[Type]) -> usize {
    types
        .iter()
        .filter(|t| matches!(t.kind_data, KindData::Struct { .. }))
        .count()
}

/// Produce C struct declarations for every recovered Go struct type.
/// Fields whose types we can't easily express in C (other structs, slices,
/// interfaces) become opaque byte arrays of the right size so the layout
/// stays correct in the disassembler's struct view.
/// C keywords and Ghidra-CParser reserved words that would collide with a
/// Go field name. Any field whose sanitized name lands in this set gets a
/// trailing underscore so the C parser doesn't reject the whole struct.
///
/// The collision is real: `runtime.boundsError` has a Go field named
/// `signed bool`, which renders as `unsigned char signed;` and breaks the
/// parse silently. Same story for any Go type with a field named after a
/// C type-specifier or storage-class keyword.
fn is_c_reserved(name: &str) -> bool {
    matches!(
        name,
        "auto"
            | "break"
            | "case"
            | "char"
            | "const"
            | "continue"
            | "default"
            | "do"
            | "double"
            | "else"
            | "enum"
            | "extern"
            | "float"
            | "for"
            | "goto"
            | "if"
            | "inline"
            | "int"
            | "long"
            | "register"
            | "restrict"
            | "return"
            | "short"
            | "signed"
            | "sizeof"
            | "static"
            | "struct"
            | "switch"
            | "typedef"
            | "union"
            | "unsigned"
            | "void"
            | "volatile"
            | "while"
            | "_Alignas"
            | "_Alignof"
            | "_Atomic"
            | "_Bool"
            | "_Complex"
            | "_Generic"
            | "_Imaginary"
            | "_Noreturn"
            | "_Static_assert"
            | "_Thread_local"
    )
}

/// Escape a field name that would otherwise collide with a C reserved
/// keyword. Suffix with `_`. The original Go field name is preserved in
/// the recovered function comment elsewhere; this is only the C-decl
/// identifier used by the disassembler's parser.
fn escape_c_field_name(name: &str) -> String {
    if is_c_reserved(name) {
        format!("{name}_")
    } else {
        name.to_string()
    }
}

fn struct_c_decls(types: &[Type]) -> Vec<String> {
    let by_addr: std::collections::HashMap<u64, &Type> =
        types.iter().map(|t| (t.addr, t)).collect();
    let mut decls = Vec::new();
    let mut emitted_names: std::collections::HashSet<String> = std::collections::HashSet::new();

    for t in types {
        let KindData::Struct { fields } = &t.kind_data else {
            continue;
        };
        let c_name = sanitize_label(&t.name);
        if c_name.is_empty() || !emitted_names.insert(c_name.clone()) {
            continue;
        }
        let mut decl = format!("struct {} {{\n", c_name);
        let mut next_offset = 0u64;
        let mut emitted = 0usize;
        for (idx, field) in fields.iter().enumerate() {
            let fsize = field_size(field, &by_addr);
            // Zero-size fields (Go marker types like sys.NotInHeap, struct{})
            // would render as `unsigned char foo[0];` which some C parsers
            // (including Ghidra's CParser on some versions) reject. Skip them.
            if fsize == 0 {
                continue;
            }
            // Pad to the field's offset if Go inserted alignment padding.
            if field.offset > next_offset {
                let pad = field.offset - next_offset;
                decl.push_str(&format!("    unsigned char _pad_{idx}[{pad}];\n"));
                next_offset = field.offset;
            }
            let (prefix, suffix) = c_type_for_field(field, &by_addr);
            let fname = sanitize_label(&field.name);
            let fname = if fname.is_empty() || fname == "_" {
                format!("_anon_{idx}")
            } else {
                escape_c_field_name(&fname)
            };
            decl.push_str(&format!("    {prefix} {fname}{suffix};\n"));
            next_offset += fsize;
            emitted += 1;
        }
        if next_offset < t.size {
            let tail = t.size - next_offset;
            decl.push_str(&format!("    unsigned char _tail[{tail}];\n"));
        }
        // Empty structs would emit `struct foo {};` which Ghidra rejects.
        if emitted == 0 && t.size == 0 {
            // Zero-size struct (Go's `struct{}` family). Skip; nothing to model.
            continue;
        }
        decl.push_str("};\n");
        decls.push(decl);
    }
    decls
}

fn field_size(field: &StructField, by_addr: &std::collections::HashMap<u64, &Type>) -> u64 {
    by_addr.get(&field.typ).map(|t| t.size).unwrap_or(8)
}

/// Map a Go field to a C type that has the right byte size. We don't try to
/// preserve semantic types (the field's "real" type is in the unstrip JSON);
/// we only need correct sizing so the disassembler's struct view shows
/// field-aligned offsets.
/// Returns `(type_prefix, suffix)` so callers can emit `<prefix> <name><suffix>`.
/// C array fields look like `unsigned char field[16];` (suffix is `[16]`);
/// most other types have empty suffix.
fn c_type_for_field(
    field: &StructField,
    by_addr: &std::collections::HashMap<u64, &Type>,
) -> (String, String) {
    let Some(t) = by_addr.get(&field.typ) else {
        return ("void *".into(), String::new());
    };
    match t.kind {
        KindName::Bool | KindName::Uint8 | KindName::Int8 => {
            ("unsigned char".into(), String::new())
        }
        KindName::Int16 | KindName::Uint16 => ("unsigned short".into(), String::new()),
        KindName::Int32 | KindName::Uint32 | KindName::Float32 => {
            ("unsigned int".into(), String::new())
        }
        KindName::Int64
        | KindName::Uint64
        | KindName::Float64
        | KindName::Int
        | KindName::Uint
        | KindName::Uintptr => ("unsigned long long".into(), String::new()),
        KindName::Pointer
        | KindName::UnsafePointer
        | KindName::Chan
        | KindName::Map
        | KindName::Func => ("void *".into(), String::new()),
        // Sized aggregates become opaque byte arrays. Disassembler still
        // sees the right struct layout, and the field name preserves intent.
        _ => ("unsigned char".into(), format!("[{}]", t.size)),
    }
}

/// Python single-quoted string literal with `\` and `'` escaped. Go symbol
/// names can contain `()`, `[]`, `*`, and other punctuation but no quotes;
/// escaping `'`, `\`, and newlines (for multi-line comments) is sufficient.
fn escape_py(s: &str) -> String {
    let mut out = String::with_capacity(s.len() + 2);
    out.push('\'');
    for c in s.chars() {
        match c {
            '\\' => out.push_str("\\\\"),
            '\'' => out.push_str("\\'"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            c if (c as u32) < 0x20 => out.push_str(&format!("\\x{:02x}", c as u32)),
            c => out.push(c),
        }
    }
    out.push('\'');
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn escapes_quotes_and_backslashes() {
        assert_eq!(escape_py("a'b"), "'a\\'b'");
        assert_eq!(escape_py("a\\b"), "'a\\\\b'");
    }

    #[test]
    fn passes_normal_go_names() {
        assert_eq!(escape_py("main.main"), "'main.main'");
        assert_eq!(escape_py("(*os.File).Read"), "'(*os.File).Read'");
    }

    #[test]
    fn sanitize_label_handles_go_punctuation() {
        assert_eq!(sanitize_label("main.main"), "main.main");
        assert_eq!(sanitize_label("(*os.File).Read"), "__os.File_.Read");
        assert_eq!(
            sanitize_label("main.Set[go.shape.string]"),
            "main.Set_go.shape.string_"
        );
        assert_eq!(sanitize_label("123main"), "_123main");
    }

    #[test]
    fn escapes_newlines_in_comments() {
        assert_eq!(escape_py("a\nb"), "'a\\nb'");
    }
}