"""Re-apply all HDCP TA annotations to psp_out/elf/hdcp.elf.

Run inside IDA (or via the ida_execute_python MCP tool) against the hdcp.elf DB.
Without these, Hex-Rays emits bare `sub_401058();` with the arguments DROPPED,
which is what made the overflow look inert for most of this project.
"""
import idc, ida_name, ida_funcs, ida_typeinf, ida_hexrays, ida_bytes, idautils
ida_hexrays.init_hexrays_plugin()

FUNCS = {
 0x401058: ("ta_memcpy",           "void *__fastcall ta_memcpy(void *dst, const void *src, unsigned int n);"),
 0x4010E0: ("ta_memcpy_aligned",   "void *__fastcall ta_memcpy_aligned(void *dst, const void *src, unsigned int n);"),
 0x401188: ("ta_memclr",           "void __fastcall ta_memclr(void *dst, unsigned int n);"),
 0x401144: ("ta_memclr_unaligned", "void __fastcall ta_memclr_unaligned(void *dst, unsigned int n);"),
 0x40182C: ("ta_bounded_copy",     "int __fastcall ta_bounded_copy(void *dst, void *src, unsigned int n);"),
 0x4045A0: ("asd_rpc_1a2",         "int __fastcall asd_rpc_1a2(unsigned int devid);"),
 0x404666: ("asd_rpc_1a3",         "int __fastcall asd_rpc_1a3(unsigned int devid, int flag);"),
 0x404646: ("asd_rpc_1a8",         "int __fastcall asd_rpc_1a8(void *buf);"),
 0x404574: ("asd_rpc_1a7",         "int __fastcall asd_rpc_1a7(unsigned int devid, void *buf, unsigned int cnt, void *out);"),
 0x40471A: ("asd_rpc_303",         "int __fastcall asd_rpc_303(void *rec5);"),
 0x4047A6: ("asd_rpc_200",         "int __fastcall asd_rpc_200(void *buf);"),
 0x404358: ("asd_topology_query",  "int __fastcall asd_topology_query(unsigned int handle, void *dst);"),
}
NAMES = {
 0x402A88: "hdcp_subcmd12_rxinfo", 0x402910: "hdcp_subcmd_dispatch",
 0x4021A0: "hdcp_cmd0F_handler",   0x4013B4: "hdcp_invoke_entry",
 0x401EBC: "hdcp_cmd0E_open",      0x401888: "hdcp_cmd01_open",
 0x404334: "__stack_chk_guard",    0x404338: "__stack_chk_fail",
}
STRUCT = """
struct hdcp_session {
  unsigned int  state;          // +0x000  2 = authenticated
  unsigned int  devid;          // +0x004
  unsigned int  active;         // +0x008  1 = slot in use
  unsigned int  f0c;
  unsigned int  topology_type;  // +0x010  from ASD 0x103; ==6 gate
  unsigned int  f14;
  unsigned int  err_state;      // +0x018  set to 7 on error
  unsigned int  f1c;
  unsigned int  f20;
  unsigned int  mode;           // +0x024
  unsigned int  f28;  unsigned int f2c;
  unsigned int  f30;  unsigned int f34;
  unsigned char rx_count;       // +0x038  ATTACKER BYTE, copy len = 5*N
  unsigned char pad39[3];
  unsigned int  f3c;  unsigned int f40;  unsigned int f44;
  unsigned int  seq_valid;      // +0x048
  unsigned int  f4c;  unsigned int f50;
  unsigned char pad54;
  unsigned char rx_ids[504];    // +0x055  SOURCE of the overflow
  unsigned char flag_24d;
  unsigned char pad24e[6];
  unsigned int  seq_num;        // +0x254
  unsigned int  hash[4];        // +0x258
  unsigned char buf268[36];     // +0x268
};
"""
def run():
    for ea,(nm,proto) in FUNCS.items():
        ida_name.set_name(ea, nm, ida_name.SN_CHECK|ida_name.SN_FORCE); idc.SetType(ea, proto)
    for ea,nm in NAMES.items():
        ida_name.set_name(ea, nm, ida_name.SN_CHECK|ida_name.SN_FORCE)
    idc.parse_decls(STRUCT, 0)
    # hdcp_verify_mac is a mid-function BL target; split it out first
    if not ida_funcs.get_func(0x403FF4):
        f = ida_funcs.get_func(0x403FD0)
        if f: ida_funcs.set_func_end(f.start_ea, 0x403FF4)
        ida_funcs.add_func(0x403FF4, 0x40423C)
    ida_name.set_name(0x403FF4, "hdcp_verify_mac", ida_name.SN_CHECK|ida_name.SN_FORCE)
    idc.SetType(0x403FF4, "int __fastcall hdcp_verify_mac(void *buf, unsigned int len, void *out32, int *outlen, unsigned int maxlen);")
    idc.SetType(0x402A88, "int __fastcall hdcp_subcmd12_rxinfo(struct hdcp_session *sess, unsigned __int8 *msg, unsigned int len, int *out);")
    for ea in (0x402A88,0x402910,0x4021A0,0x401EBC): ida_hexrays.mark_cfunc_dirty(ea)
    return "annotations applied"
run()