// Frida hooks for the 32-bit MIDL COM proxy path used by NmxSvcps.dll.
// Usage:
//   frida -f src/MxTraceHarness/bin/Release/net481/MxTraceHarness.exe \
//     -l analysis/frida/nmx-com-proxy-trace.js -- --scenario=register --duration=3

const maxDump = 8192;
const installed = {};
let lastBstrMarshalBuffer = ptr("0");
let lastBstrMarshalLength = 0;

function now() {
  return new Date().toISOString();
}

function emit(event) {
  event.time = now();
  console.log(JSON.stringify(event));
}

function toHex(arrayBuffer) {
  if (arrayBuffer === null || arrayBuffer === "") return "";
  const bytes = new Uint8Array(arrayBuffer);
  const out = [];
  for (let i = 0; i < bytes.length; i++) out.push(bytes[i].toString(16).padStart(2, "0"));
  return out.join(" ");
}

function dump(ptrValue, length) {
  try {
    if (ptrValue.isNull() || length <= 0) return "";
    return toHex(ptrValue.readByteArray(Math.min(length, maxDump)));
  } catch (e) {
    return null;
  }
}

function readU32(ptrValue, offset) {
  try {
    return ptrValue.add(offset).readU32();
  } catch (e) {
    return null;
  }
}

function readPtr(ptrValue, offset) {
  try {
    return ptrValue.add(offset).readPointer();
  } catch (e) {
    return ptr("0");
  }
}

function moduleNameFor(address) {
  try {
    const module = Process.findModuleByAddress(address);
    return module === null ? "" : module.name;
  } catch (e) {
    return "";
  }
}

function readUtf16(ptrValue, maxChars) {
  try {
    if (ptrValue.isNull()) return "";
    return ptrValue.readUtf16String(maxChars);
  } catch (e) {
    return "";
  }
}

function describePossibleBstr(ptrValue) {
  try {
    if (ptrValue.isNull()) return null;
    const byteLength = ptrValue.sub(4).readU32();
    if (byteLength > 4096 || (byteLength % 2) !== 0) return null;
    return {
      ptr: ptrValue.toString(),
      byteLength,
      charLength: byteLength / 2,
      text: readUtf16(ptrValue, byteLength / 2)
    };
  } catch (e) {
    return null;
  }
}

function dumpNdrVarArgs(args) {
  const values = [];
  for (let i = 2; i < 12; i++) {
    let value = ptr("0");
    try {
      value = args[i];
    } catch (e) {
      value = ptr("0");
    }

    const item = {
      index: i,
      value: value.toString(),
      module: moduleNameFor(value),
      asU32: value.toUInt32()
    };

    const bstr = describePossibleBstr(value);
    if (bstr !== null) {
      item.bstr = bstr;
    }

    values.push(item);
  }

  return values;
}

function dumpStackWords(context, count) {
  const values = [];
  let esp = ptr("0");
  try {
    esp = context.esp;
  } catch (e) {
    return values;
  }

  for (let i = 0; i < count; i++) {
    const address = esp.add(i * Process.pointerSize);
    let value = ptr("0");
    try {
      value = address.readPointer();
    } catch (e) {
      value = ptr("0");
    }

    const item = {
      index: i,
      address: address.toString(),
      value: value.toString(),
      module: moduleNameFor(value),
      asU32: value.toUInt32()
    };

    const bstr = describePossibleBstr(value);
    if (bstr !== null) {
      item.bstr = bstr;
    }

    values.push(item);
  }

  return values;
}

function hookExport(moduleName, exportName, callbacks) {
  const key = moduleName + "!" + exportName;
  if (installed[key]) return;

  let address = null;
  try {
    const module = Process.findModuleByName(moduleName);
    if (module !== null && typeof module.findExportByName === "function") {
      address = module.findExportByName(exportName);
    }
    if (address === null && typeof Module.findExportByName === "function") {
      address = Module.findExportByName(moduleName, exportName);
    }
  } catch (e) {
    address = null;
  }

  if (address === null) {
    emit({ event: "hook.missing", module: moduleName, name: exportName });
    return;
  }

  Interceptor.attach(address, callbacks(address));
  installed[key] = true;
  emit({ event: "hook.installed", module: moduleName, name: exportName, address: address.toString() });
}

function dumpRpcMessage(pRpcMsg) {
  // 32-bit RPC_MESSAGE layout:
  // Handle, DataRepresentation, Buffer, BufferLength, ProcNum, TransferSyntax,
  // RpcInterfaceInformation, ReservedForRuntime, ManagerEpv, ImportContext, RpcFlags.
  const buffer = readPtr(pRpcMsg, 8);
  const bufferLength = readU32(pRpcMsg, 12);
  const procNum = readU32(pRpcMsg, 16);
  const rpcInterfaceInfo = readPtr(pRpcMsg, 24);

  return {
    rpcMessage: pRpcMsg.toString(),
    buffer: buffer.toString(),
    bufferLength,
    procNum,
    rpcInterfaceInfo: rpcInterfaceInfo.toString(),
    rpcInterfaceModule: moduleNameFor(rpcInterfaceInfo),
    hex: bufferLength === null ? null : dump(buffer, bufferLength)
  };
}

function dumpStubMessage(pStubMsg) {
  // Common 32-bit MIDL_STUB_MESSAGE prefix. This is intentionally broad; the
  // important fields for this investigation are the RPC_MESSAGE pointer and
  // the active buffer range used by the generated NmxSvcps proxy.
  const rpcMsg = readPtr(pStubMsg, 0);
  const buffer = readPtr(pStubMsg, 4);
  const bufferStart = readPtr(pStubMsg, 8);
  const bufferEnd = readPtr(pStubMsg, 12);
  const bufferLength = readU32(pStubMsg, 20);
  let activeLength = 0;
  try {
    activeLength = bufferEnd.sub(bufferStart).toInt32();
  } catch (e) {
    activeLength = 0;
  }

  return {
    stubMessage: pStubMsg.toString(),
    rpcMessage: rpcMsg.toString(),
    buffer: buffer.toString(),
    bufferStart: bufferStart.toString(),
    bufferEnd: bufferEnd.toString(),
    bufferLength,
    activeLength,
    bufferHex: activeLength > 0 ? dump(bufferStart, activeLength) : "",
    rpc: rpcMsg.isNull() ? null : dumpRpcMessage(rpcMsg)
  };
}

function installHooks() {
  hookExport("oleaut32.dll", "BSTR_UserMarshal", function () {
    return {
      onEnter(args) {
        this.buffer = args[1];
        this.bstrSlot = args[2];
        let bstr = ptr("0");
        try {
          bstr = this.bstrSlot.readPointer();
        } catch (e) {
          bstr = ptr("0");
        }
        emit({
          event: "bstr.usermarshal.enter",
          callerModule: moduleNameFor(this.returnAddress),
          flags: args[0].toString(),
          buffer: this.buffer.toString(),
          bstrSlot: this.bstrSlot.toString(),
          bstr: describePossibleBstr(bstr)
        });
      },
      onLeave(retval) {
        let length = 0;
        try {
          length = retval.sub(this.buffer).toInt32();
        } catch (e) {
          length = 0;
        }
        lastBstrMarshalBuffer = this.buffer;
        lastBstrMarshalLength = length;
        emit({
          event: "bstr.usermarshal.leave",
          callerModule: moduleNameFor(this.returnAddress),
          buffer: this.buffer ? this.buffer.toString() : "",
          retval: retval.toString(),
          marshaledLength: length,
          marshaledHex: length > 0 ? dump(this.buffer, length) : ""
        });
      }
    };
  });

  hookExport("rpcrt4.dll", "NdrInterfacePointerMarshall", function () {
    return {
      onEnter(args) {
        this.stubMsg = args[0];
        emit({
          event: "ndr.interfaceptr.marshal.enter",
          callerModule: moduleNameFor(this.returnAddress),
          interfacePointer: args[1].toString(),
          format: args[2].toString(),
          formatPrefix: dump(args[2], 32),
          stub: dumpStubMessage(this.stubMsg)
        });
      },
      onLeave(retval) {
        emit({
          event: "ndr.interfaceptr.marshal.leave",
          callerModule: moduleNameFor(this.returnAddress),
          retval: retval.toString(),
          stub: this.stubMsg ? dumpStubMessage(this.stubMsg) : null
        });
      }
    };
  });

  hookExport("rpcrt4.dll", "NdrClientCall2", function () {
    return {
      onEnter(args) {
        this.format = args[1];
        this.varargs = dumpNdrVarArgs(args);
        emit({
          event: "ndr.client.enter",
          callerModule: moduleNameFor(this.returnAddress),
          stubDesc: args[0].toString(),
          procFormat: this.format.toString(),
          procFormatPrefix: dump(this.format, 64),
          varargs: this.varargs,
          stack: dumpStackWords(this.context, 24)
        });
      },
      onLeave(retval) {
        let surroundingStubHex = "";
        if (lastBstrMarshalBuffer && !lastBstrMarshalBuffer.isNull()) {
          try {
            surroundingStubHex = dump(lastBstrMarshalBuffer.sub(36), 192);
          } catch (e) {
            surroundingStubHex = null;
          }
        }
        emit({
          event: "ndr.client.leave",
          callerModule: moduleNameFor(this.returnAddress),
          procFormat: this.format ? this.format.toString() : "",
          retval: retval.toString(),
          lastBstrMarshalLength,
          surroundingStubHex
        });
      }
    };
  });

  hookExport("rpcrt4.dll", "NdrProxySendReceive", function () {
    return {
      onEnter(args) {
        emit({
          event: "ndr.proxy.sendreceive.enter",
          callerModule: moduleNameFor(this.returnAddress),
          thisPtr: args[0].toString(),
          stub: dumpStubMessage(args[1])
        });
      },
      onLeave(retval) {
        emit({
          event: "ndr.proxy.sendreceive.leave",
          callerModule: moduleNameFor(this.returnAddress),
          retval: retval.toString()
        });
      }
    };
  });

  hookExport("rpcrt4.dll", "I_RpcSendReceive", function () {
    return {
      onEnter(args) {
        emit({
          event: "rpc.sendreceive.enter",
          callerModule: moduleNameFor(this.returnAddress),
          rpc: dumpRpcMessage(args[0])
        });
      },
      onLeave(retval) {
        emit({
          event: "rpc.sendreceive.leave",
          callerModule: moduleNameFor(this.returnAddress),
          retval: retval.toString()
        });
      }
    };
  });
}

emit({ event: "script.loaded", process: Process.id, arch: Process.arch, pointerSize: Process.pointerSize });
installHooks();
const retryTimer = setInterval(function () {
  installHooks();
  if (installed["oleaut32.dll!BSTR_UserMarshal"]
      && installed["rpcrt4.dll!NdrInterfacePointerMarshall"]
      && installed["rpcrt4.dll!NdrClientCall2"]) {
    clearInterval(retryTimer);
  }
}, 100);