#!/usr/bin/env python3
"""Emulate the IronWorm Linux payload under Qiling to catch the AES key at
runtime, without ever touching the real network or filesystem.

Never executes the sample on real hardware/OS: Qiling interprets its
instructions in an emulated CPU + virtual filesystem rootfs. All network
syscalls are overridden below to log-and-deny rather than pass through to
the host, which is Qiling's default (unsafe) behavior for AF_INET sockets.
"""
import sys

from qiling import Qiling
from qiling.const import QL_INTERCEPT
from qiling.os.mapper import QlFsMapper

TARGET = "binaries/payload_linux"
ROOTFS = "qiling_rootfs/x8664_linux"

AES_OPCODES = {
    b"\x66\x0f\x38\xdc": "aesenc",
    b"\x66\x0f\x38\xdd": "aesenclast",
    b"\x66\x0f\x38\xde": "aesdec",
    b"\x66\x0f\x38\xdf": "aesdeclast",
    b"\x66\x0f\x3a\xdf": "aeskeygenassist",
}

seen_aes_hits = 0
MAX_AES_HITS_LOGGED = 20


def deny_network(name):
    def handler(ql, *args):
        ql.log.warning(f"[BLOCKED] {name}({args}) -- denied, no real network access")
        return -1
    return handler


def on_code(ql, address, size):
    global seen_aes_hits
    if size < 4:
        return
    try:
        data = ql.mem.read(address, 4)
    except Exception:
        return
    for opcode, name in AES_OPCODES.items():
        if data == opcode:
            seen_aes_hits += 1
            if seen_aes_hits <= MAX_AES_HITS_LOGGED:
                regs = ql.arch.regs
                print(f"[AES] {name} @ {address:#x} "
                      f"rdi={regs.rdi:#x} rsi={regs.rsi:#x} rdx={regs.rdx:#x} rcx={regs.rcx:#x}")
                for reg_name, val in (("rdi", regs.rdi), ("rsi", regs.rsi), ("rdx", regs.rdx)):
                    try:
                        blob = ql.mem.read(val, 32)
                        print(f"    [{reg_name}={val:#x}] -> {blob.hex()}")
                    except Exception:
                        pass
            break


def main():
    ql = Qiling([TARGET], ROOTFS, verbose=1)

    for syscall_name in ("connect", "sendto", "send", "bind", "recvfrom", "recv"):
        ql.os.set_api(syscall_name, deny_network(syscall_name), QL_INTERCEPT.CALL)

    ql.hook_code(on_code)

    try:
        ql.run(timeout=int(60 * 1e6))  # microseconds; hard cap so a hang doesn't run forever
    except Exception as e:
        print(f"[emulation stopped] {e}")

    print(f"\nTotal AES instruction hits observed: {seen_aes_hits}")


if __name__ == "__main__":
    main()