'use strict';

const moduleName = 'aahClientManaged.dll';
const knownModulePaths = [
  moduleName,
  'C:\\Users\\dohertj2\\Desktop\\histsdk\\current\\aahClientManaged.dll'
];
const maxDumpBytes = 256;

const targets = [
  ['CServiceUtility.SaveOpenConnectionParams', 0x2f0d54],
  ['CClientInfo.SerializeOpenConnectionInParams', 0x2f0928],
  ['CClientInfo.SerializeOpenConnectionInParams2', 0x2f0690],
  ['CClientInfo.SerializeOpenConnectionInParams2Content', 0x2f03a4],
  ['CClientInfo.SerializeOpenConnectionInParams3', 0x2f00c0],
  ['CClientInfo.SerializeOpenConnectionInParams3Content', 0x2efdac],
  ['CClientInfo.SerializeOpenConnectionInParams4', 0x2f2e88],
  ['CClientInfo.EncryptWithClientKey', 0x0],
  ['CHistoryConnectionWCF.GetClientKey', 0x2f3dc4],
  ['CHistoryConnectionWCF.OpenConnection', 0x2ec9c8],
  ['CHistoryConnectionWCF.OpenConnection2', 0x2fdeb8],
  ['CHistoryConnectionWCF.OpenConnection3', 0x2fedb4],
  ['CHistoryConnectionWCF.RegisterTags', 0x2f6f78],
  ['CHistoryConnectionWCF.ValidateClientCredential', 0x302e90],
  ['HistorianClient.OpenConnection', 0x4170e8],
  ['HistorianAccess.AddTagInternal', 0x43be68],
  ['HistorianAccess.CreateDefaultEventTag', 0x43c2d4],
  ['HistorianClient.AddHistorianTag', 0x417c18],
  ['HistorianClient.ConvertEventTagToTagMetadata', 0x417b68],
  ['HistorianClient.StartQuery', 0x415bbc],
  ['HistorianClient.StartEventQuery', 0x41811c],
  ['HistorianClient.StartDataQuery', 0x4160c4],
  ['CTagMetadata.Save<SByteStream<SCrtMemFile>>', 0x1044dc],
  ['ClientApp.StartDataQuery', 0x400f9c],
  ['ClientApp.StartEventQuery', 0x4015a4],
  ['Query.StartDataQuery', 0x41cacc],
  ['CRetrievalConnectionWCF.StartQuery2', 0x36eb48],
  ['CRetrievalConnectionWCF.StartEventQuery', 0x370324],
  ['HistoryQuery.StartQuery', 0x44012c],
  ['EventQuery.StartQuery', 0x43035c],
  ['Query.StartEventQuery', 0x41db4c],
  ['EventQueryFilters.Save<SCrtMemFile>', 0x41d38c],
  ['EventQueryRequest.Save<SCrtMemFile>', 0x41d48c],
  ['QueryColumnSelector.ctor.default', 0x1b94c],
  ['QueryColumnSelector.SelectNonSummaryColumns', 0x1ee34],
  ['QueryColumnSelector.Save<SCrtMemFile>', 0x41b8d8],
  ['QueryColumnSelector.GetColumnSelectorFlags', 0x41e110],
  ['HistorianClient.GetNextRow<DataQueryResultRow>', 0x42f818],
  ['DataQueryResultBuffer.GetNextRow<SByteStream<SCrtMemFile>>', 0x42f6f8],
  ['Query.GetNextRow', 0x42f744],
  ['HistorianClient.GetNextRow<EventQueryResultRow>', 0x430e10],
  ['EventQueryResultBuffer.GetNextRow<SByteStream<SCrtMemFile>>', 0x430a3c],
  ['Event.Query.GetNextRow', 0x430af4]
].filter(t => t[1] !== 0);

let startupLogged = false;
let hooksInstalled = false;
let getModuleHandleW = null;

function emit(kind, payload) {
  payload.kind = kind;
  payload.pid = Process.id;
  payload.tid = Process.getCurrentThreadId();
  payload.timestamp = new Date().toISOString();
  console.log('FRIDA_EVENT ' + JSON.stringify(payload));
}

function isReadablePointer(value) {
  if (value.isNull()) {
    return false;
  }

  if (value.compare(ptr('0x10000')) < 0) {
    return false;
  }

  const range = Process.findRangeByAddress(value);
  return range !== null && range.protection.indexOf('r') !== -1;
}

function toHex(bytes) {
  const parts = [];
  for (let i = 0; i < bytes.length; i++) {
    const text = bytes[i].toString(16);
    parts.push(text.length === 1 ? '0' + text : text);
  }
  return parts.join('');
}

function asciiPreview(bytes) {
  let text = '';
  for (let i = 0; i < bytes.length; i++) {
    const b = bytes[i];
    text += b >= 32 && b <= 126 ? String.fromCharCode(b) : '.';
  }
  return text;
}

function utf16Preview(bytes) {
  let text = '';
  const count = Math.min(bytes.length - (bytes.length % 2), 96);
  for (let i = 0; i < count; i += 2) {
    const code = bytes[i] | (bytes[i + 1] << 8);
    text += code >= 32 && code <= 126 ? String.fromCharCode(code) : '.';
  }
  return text;
}

function dumpPointer(value) {
  const range = Process.findRangeByAddress(value);
  const available = range === null ? 0 : Math.min(maxDumpBytes, range.base.add(range.size).sub(value).toNumber());
  if (available <= 0) {
    return null;
  }

  const raw = Memory.readByteArray(value, available);
  if (raw === null) {
    return null;
  }

  const bytes = Array.from(new Uint8Array(raw));
  return {
    address: value.toString(),
    rangeBase: range.base.toString(),
    rangeSize: range.size,
    protection: range.protection,
    byteCount: bytes.length,
    hexPrefix: toHex(bytes.slice(0, Math.min(bytes.length, 96))),
    asciiPrefix: asciiPreview(bytes.slice(0, Math.min(bytes.length, 96))),
    utf16Prefix: utf16Preview(bytes.slice(0, Math.min(bytes.length, 192)))
  };
}

function inspectArgs(args, count) {
  const inspected = [];
  for (let i = 0; i < count; i++) {
    const value = args[i];
    const item = {
      index: i,
      value: value.toString()
    };

    if (isReadablePointer(value)) {
      try {
        item.memory = dumpPointer(value);
      } catch (e) {
        item.memoryError = String(e);
      }
    }

    inspected.push(item);
  }

  return inspected;
}

function hookTarget(base, name, rva) {
  const address = base.add(rva);
  try {
    Interceptor.attach(address, {
      onEnter(args) {
        this.name = name;
        this.argsSnapshot = inspectArgs(args, 10);
        emit('enter', {
          function: name,
          rva: '0x' + rva.toString(16),
          address: address.toString(),
          args: this.argsSnapshot
        });
      },
      onLeave(retval) {
        emit('leave', {
          function: this.name,
          retval: retval.toString(),
          args: this.argsSnapshot
        });
      }
    });
    emit('hooked', { function: name, rva: '0x' + rva.toString(16), address: address.toString() });
  } catch (e) {
    emit('hook-error', { function: name, rva: '0x' + rva.toString(16), address: address.toString(), error: String(e) });
  }
}

function installHooks() {
  if (!startupLogged) {
    startupLogged = true;
    emit('startup', {
      arch: Process.arch,
      platform: Process.platform,
      modules: Process.enumerateModules()
        .filter(m => m.name.toLowerCase().indexOf('aah') !== -1 || m.path.toLowerCase().indexOf('histsdk') !== -1)
        .map(m => ({ name: m.name, base: m.base.toString(), size: m.size, path: m.path }))
    });
  }

  const modules = Process.enumerateModules().filter(m => m.name.toLowerCase() === moduleName.toLowerCase());
  let module = modules.length > 0 ? modules[0] : null;
  let base = module === null ? null : module.base;

  if (base === null && getModuleHandleW !== null) {
    for (const candidate of knownModulePaths) {
      const handle = getModuleHandleW(Memory.allocUtf16String(candidate));
      if (!handle.isNull()) {
        base = handle;
        emit('module-handle-found', { module: candidate, base: base.toString() });
        break;
      }
    }
  }

  if (base === null) {
    return false;
  }

  if (hooksInstalled) {
    return true;
  }

  hooksInstalled = true;
  emit('module-loaded', { module: moduleName, base: base.toString(), arch: Process.arch, platform: Process.platform });
  for (const [name, rva] of targets) {
    hookTarget(base, name, rva);
  }

  return true;
}

function isInterestingModule(module) {
  const name = module.name.toLowerCase();
  const path = module.path.toLowerCase();
  return name.indexOf('aah') !== -1
    || name.indexOf('historian') !== -1
    || path.indexOf('histsdk') !== -1
    || path.indexOf('aveva') !== -1;
}

try {
  Process.attachModuleObserver({
    onAdded(module) {
      if (isInterestingModule(module)) {
        emit('module-added', {
          name: module.name,
          base: module.base.toString(),
          size: module.size,
          path: module.path
        });
      }

      installHooks();
    }
  });
} catch (e) {
  emit('module-observer-error', { error: String(e) });
}

try {
  const kernel32 = Process.getModuleByName('kernel32.dll');
  getModuleHandleW = new NativeFunction(kernel32.getExportByName('GetModuleHandleW'), 'pointer', ['pointer']);
  const hookLoader = (exportName, pathArgIndex) => {
    const fn = kernel32.getExportByName(exportName);
    Interceptor.attach(fn, {
      onEnter(args) {
        this.exportName = exportName;
        this.path = args[pathArgIndex].isNull() ? '' : args[pathArgIndex].readUtf16String();
      },
      onLeave(retval) {
        const lower = (this.path || '').toLowerCase();
        if (lower.indexOf('aah') !== -1 || lower.indexOf('historian') !== -1 || lower.indexOf('histsdk') !== -1 || lower.indexOf('aveva') !== -1) {
          emit('load-library', { api: this.exportName, path: this.path, result: retval.toString() });
        }
        installHooks();
      }
    });
  };

  hookLoader('LoadLibraryW', 0);
  hookLoader('LoadLibraryExW', 0);
} catch (e) {
  emit('load-library-hook-error', { error: String(e) });
}

try {
  const ntdll = Process.getModuleByName('ntdll.dll');
  const ldrLoadDll = ntdll.getExportByName('LdrLoadDll');
  Interceptor.attach(ldrLoadDll, {
    onEnter(args) {
      this.path = '';
      try {
        const unicodeString = args[2];
        const length = unicodeString.readU16();
        const buffer = unicodeString.add(Process.pointerSize * 2).readPointer();
        if (!buffer.isNull() && length > 0) {
          this.path = buffer.readUtf16String(length / 2);
        }
      } catch (e) {
        this.path = '<read-error:' + String(e) + '>';
      }
    },
    onLeave(retval) {
      const lower = (this.path || '').toLowerCase();
      if (lower.indexOf('aah') !== -1 || lower.indexOf('historian') !== -1 || lower.indexOf('histsdk') !== -1 || lower.indexOf('aveva') !== -1) {
        emit('ldr-load-dll', { path: this.path, result: retval.toString() });
      }
      installHooks();
    }
  });
} catch (e) {
  emit('ldr-load-dll-hook-error', { error: String(e) });
}

if (!installHooks()) {
  const timer = setInterval(() => {
    if (installHooks()) {
      clearInterval(timer);
    }
  }, 50);
}