from __future__ import annotations

import argparse
import socket
import struct
import uuid


IOBJECT_EXPORTER = uuid.UUID("99fcfec4-5260-101b-bbcb-00aa0021347a")
NDR20 = uuid.UUID("8a885d04-1ceb-11c9-9fe8-08002b104860")


def uuid_le(value: uuid.UUID) -> bytes:
    return value.bytes_le


def pdu_header(ptype: int, flags: int, frag_len: int, call_id: int) -> bytes:
    return struct.pack("<BBBBLHHI", 5, 0, ptype, flags, 0x10, frag_len, 0, call_id)


def bind_pdu(call_id: int = 1) -> bytes:
    body = struct.pack("<HHIBBBB", 4280, 4280, 0, 1, 0, 0, 0)
    body += struct.pack("<HBB", 0, 1, 0)
    body += uuid_le(IOBJECT_EXPORTER) + struct.pack("<HH", 0, 0)
    body += uuid_le(NDR20) + struct.pack("<HH", 2, 0)
    return pdu_header(11, 3, 16 + len(body), call_id) + body


def request_pdu(call_id: int, opnum: int, stub: bytes) -> bytes:
    body = struct.pack("<IHH", len(stub), 0, opnum) + stub
    return pdu_header(0, 3, 16 + len(body), call_id) + body


def encode_resolve_oxid_stub(oxid: int, protseqs: list[int]) -> bytes:
    # IObjectExporter::ResolveOxid, opnum 0:
    # OXID pOxid; unsigned short cRequestedProtseqs;
    # conformant ushort array arRequestedProtseqs.
    body = struct.pack("<QH", oxid, len(protseqs))
    body += b"\x00\x00"
    body += struct.pack("<I", len(protseqs))
    body += b"".join(struct.pack("<H", protseq) for protseq in protseqs)
    if len(body) % 4:
        body += b"\x00" * (4 - len(body) % 4)
    return body


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=135)
    parser.add_argument("--oxid", required=True, help="hex OXID, for example 0xEAF0D2B53BAB5BC2")
    parser.add_argument("--protseq", type=lambda text: int(text, 0), action="append", default=[7])
    args = parser.parse_args()

    oxid = int(args.oxid, 0)
    stub = encode_resolve_oxid_stub(oxid, args.protseq)
    print(f"resolve_oxid_stub={stub.hex(' ')}")

    with socket.create_connection((args.host, args.port), timeout=5) as sock:
        sock.sendall(bind_pdu())
        bind_response = sock.recv(4096)
        print(f"bind_response={bind_response.hex(' ')}")

        sock.sendall(request_pdu(2, 0, stub))
        response = sock.recv(8192)
        print(f"resolve_response={response.hex(' ')}")

        if len(response) >= 32 and response[2] == 2:
            stub_data = response[24:]
            print(f"resolve_response_stub={stub_data.hex(' ')}")
            if len(stub_data) >= 4:
                print(f"likely_error_status=0x{int.from_bytes(stub_data[-4:], 'little'):08x}")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())