'use strict';

const historianPort = 32568;
const maxPrefixBytes = 16;
const sockets = new Map();
const handles = new Map();
let getPeerName = null;

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

function socketKey(socket) {
  return socket.toString();
}

function readPortNetworkOrder(address) {
  const hi = address.readU8();
  const lo = address.add(1).readU8();
  return (hi << 8) | lo;
}

function parseSockaddr(address) {
  if (address.isNull()) {
    return null;
  }

  try {
    const family = address.readU16();
    if (family === 2) {
      const port = readPortNetworkOrder(address.add(2));
      const a = [
        address.add(4).readU8(),
        address.add(5).readU8(),
        address.add(6).readU8(),
        address.add(7).readU8()
      ].join('.');
      return { family: 'IPv4', address: a, port };
    }

    if (family === 23) {
      const port = readPortNetworkOrder(address.add(2));
      const parts = [];
      for (let i = 0; i < 16; i += 2) {
        const value = (address.add(8 + i).readU8() << 8) | address.add(8 + i + 1).readU8();
        parts.push(value.toString(16));
      }
      return { family: 'IPv6', address: parts.join(':'), port };
    }

    return { family: String(family), address: null, port: null };
  } catch (e) {
    return { error: String(e) };
  }
}

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 payloadSummary(buffer, length) {
  const byteCount = Math.max(0, Math.min(length, maxPrefixBytes));
  if (buffer.isNull() || byteCount === 0) {
    return { byteCount: length, prefixByteCount: 0, prefixHex: '' };
  }

  try {
    const raw = Memory.readByteArray(buffer, byteCount);
    return {
      byteCount: length,
      prefixByteCount: byteCount,
      prefixHex: toHex(Array.from(new Uint8Array(raw)))
    };
  } catch (e) {
    return { byteCount: length, prefixByteCount: 0, prefixHex: '', error: String(e) };
  }
}

function isHistorianSocket(socket) {
  ensurePeer(socket);
  const peer = sockets.get(socketKey(socket));
  return peer !== undefined && peer.port === historianPort;
}

function ensurePeer(socket) {
  const key = socketKey(socket);
  if (sockets.has(key) || getPeerName === null) {
    return;
  }

  try {
    const address = Memory.alloc(128);
    const length = Memory.alloc(4);
    length.writeU32(128);
    const result = getPeerName(socket, address, length);
    if (result === 0) {
      const peer = parseSockaddr(address);
      if (peer !== null) {
        sockets.set(key, peer);
        if (peer.port === historianPort) {
          emit('peer-discovered', { socket: key, peer });
        }
      }
    }
  } catch (e) {
    emit('peer-discovery-error', { socket: key, error: String(e) });
  }
}

function hookExport(moduleName, exportName, callbacks) {
  try {
    const module = Process.getModuleByName(moduleName);
    const address = module.getExportByName(exportName);
    Interceptor.attach(address, callbacks);
    emit('hooked', { api: exportName, module: moduleName, address: address.toString() });
  } catch (e) {
    emit('hook-error', { api: exportName, module: moduleName, error: String(e) });
  }
}

function isInterestingPath(path) {
  const lower = (path || '').toLowerCase();
  return lower.indexOf('\\pipe\\') !== -1
    || lower.indexOf('historian') !== -1
    || lower.indexOf('aveva') !== -1
    || lower.indexOf('wonderware') !== -1
    || lower.indexOf('aah') !== -1
    || lower.indexOf('ww') !== -1;
}

function readUnicodeString(unicodeString) {
  if (unicodeString.isNull()) {
    return '';
  }

  try {
    const length = unicodeString.readU16();
    const bufferOffset = Process.pointerSize === 8 ? 8 : 4;
    const buffer = unicodeString.add(bufferOffset).readPointer();
    if (buffer.isNull() || length === 0) {
      return '';
    }

    return buffer.readUtf16String(length / 2);
  } catch (e) {
    return '<unicode-read-error:' + String(e) + '>';
  }
}

function readObjectAttributesName(objectAttributes) {
  if (objectAttributes.isNull()) {
    return '';
  }

  try {
    const objectNameOffset = Process.pointerSize === 8 ? 16 : 8;
    const unicodeString = objectAttributes.add(objectNameOffset).readPointer();
    return readUnicodeString(unicodeString);
  } catch (e) {
    return '<object-attributes-read-error:' + String(e) + '>';
  }
}

emit('startup', { arch: Process.arch, platform: Process.platform, historianPort });

try {
  const ws2 = Process.getModuleByName('ws2_32.dll');
  getPeerName = new NativeFunction(ws2.getExportByName('getpeername'), 'int', ['pointer', 'pointer', 'pointer']);
} catch (e) {
  emit('getpeername-error', { error: String(e) });
}

hookExport('ws2_32.dll', 'connect', {
  onEnter(args) {
    this.socket = args[0];
    this.peer = parseSockaddr(args[1]);
  },
  onLeave(retval) {
    if (retval.toInt32() === 0 && this.peer !== null) {
      sockets.set(socketKey(this.socket), this.peer);
    }

    if (this.peer !== null && this.peer.port === historianPort) {
      emit('connect', {
        socket: socketKey(this.socket),
        peer: this.peer,
        retval: retval.toInt32()
      });
    }
  }
});

hookExport('ws2_32.dll', 'WSAConnect', {
  onEnter(args) {
    this.socket = args[0];
    this.peer = parseSockaddr(args[1]);
  },
  onLeave(retval) {
    if (retval.toInt32() === 0 && this.peer !== null) {
      sockets.set(socketKey(this.socket), this.peer);
    }

    if (this.peer !== null && this.peer.port === historianPort) {
      emit('connect', {
        api: 'WSAConnect',
        socket: socketKey(this.socket),
        peer: this.peer,
        retval: retval.toInt32()
      });
    }
  }
});

hookExport('ws2_32.dll', 'closesocket', {
  onEnter(args) {
    this.socket = args[0];
    this.peer = sockets.get(socketKey(this.socket));
  },
  onLeave(retval) {
    if (this.peer !== undefined && this.peer.port === historianPort) {
      emit('close', { socket: socketKey(this.socket), peer: this.peer, retval: retval.toInt32() });
    }

    sockets.delete(socketKey(this.socket));
  }
});

hookExport('ws2_32.dll', 'send', {
  onEnter(args) {
    this.socket = args[0];
    this.buffer = args[1];
    this.length = args[2].toInt32();
  },
  onLeave(retval) {
    if (isHistorianSocket(this.socket)) {
      emit('send', {
        api: 'send',
        socket: socketKey(this.socket),
        peer: sockets.get(socketKey(this.socket)),
        requestedBytes: this.length,
        resultBytes: retval.toInt32(),
        payload: payloadSummary(this.buffer, Math.max(0, Math.min(this.length, retval.toInt32())))
      });
    }
  }
});

hookExport('ws2_32.dll', 'recv', {
  onEnter(args) {
    this.socket = args[0];
    this.buffer = args[1];
  },
  onLeave(retval) {
    const length = retval.toInt32();
    if (length > 0 && isHistorianSocket(this.socket)) {
      emit('recv', {
        api: 'recv',
        socket: socketKey(this.socket),
        peer: sockets.get(socketKey(this.socket)),
        resultBytes: length,
        payload: payloadSummary(this.buffer, length)
      });
    }
  }
});

function readWsaBuffers(lpBuffers, count) {
  const result = [];
  const stride = Process.pointerSize === 8 ? 16 : 8;
  const pointerOffset = Process.pointerSize === 8 ? 8 : 4;
  for (let i = 0; i < count; i++) {
    const item = lpBuffers.add(i * stride);
    const length = item.readU32();
    const buffer = item.add(pointerOffset).readPointer();
    result.push(payloadSummary(buffer, length));
  }

  return result;
}

hookExport('ws2_32.dll', 'WSASend', {
  onEnter(args) {
    this.socket = args[0];
    this.buffers = args[1];
    this.count = args[2].toInt32();
    this.sentPointer = args[3];
    this.summaries = isHistorianSocket(this.socket) ? readWsaBuffers(this.buffers, this.count) : [];
  },
  onLeave(retval) {
    if (isHistorianSocket(this.socket)) {
      let sent = null;
      try {
        sent = this.sentPointer.isNull() ? null : this.sentPointer.readU32();
      } catch (_) {
        sent = null;
      }

      emit('send', {
        api: 'WSASend',
        socket: socketKey(this.socket),
        peer: sockets.get(socketKey(this.socket)),
        retval: retval.toInt32(),
        resultBytes: sent,
        buffers: this.summaries
      });
    }
  }
});

hookExport('ws2_32.dll', 'WSARecv', {
  onEnter(args) {
    this.socket = args[0];
    this.buffers = args[1];
    this.count = args[2].toInt32();
    this.receivedPointer = args[3];
  },
  onLeave(retval) {
    if (isHistorianSocket(this.socket)) {
      let received = null;
      try {
        received = this.receivedPointer.isNull() ? null : this.receivedPointer.readU32();
      } catch (_) {
        received = null;
      }

      emit('recv', {
        api: 'WSARecv',
        socket: socketKey(this.socket),
        peer: sockets.get(socketKey(this.socket)),
        retval: retval.toInt32(),
        resultBytes: received,
        buffers: readWsaBuffers(this.buffers, this.count)
      });
    }
  }
});

hookExport('kernel32.dll', 'CreateFileW', {
  onEnter(args) {
    this.path = args[0].isNull() ? '' : args[0].readUtf16String();
    this.interesting = isInterestingPath(this.path);
  },
  onLeave(retval) {
    if (this.interesting && !retval.equals(ptr('-1'))) {
      handles.set(retval.toString(), this.path);
      emit('ipc-open', {
        api: 'CreateFileW',
        handle: retval.toString(),
        path: this.path
      });
    }
  }
});

hookExport('kernel32.dll', 'CloseHandle', {
  onEnter(args) {
    this.handle = args[0];
    this.path = handles.get(this.handle.toString());
  },
  onLeave(retval) {
    if (this.path !== undefined) {
      emit('ipc-close', {
        api: 'CloseHandle',
        handle: this.handle.toString(),
        path: this.path,
        retval: retval.toInt32()
      });
      handles.delete(this.handle.toString());
    }
  }
});

hookExport('kernel32.dll', 'WriteFile', {
  onEnter(args) {
    this.handle = args[0];
    this.path = handles.get(this.handle.toString());
    this.buffer = args[1];
    this.length = args[2].toInt32();
    this.writtenPointer = args[3];
  },
  onLeave(retval) {
    if (this.path !== undefined) {
      let written = null;
      try {
        written = this.writtenPointer.isNull() ? null : this.writtenPointer.readU32();
      } catch (_) {
        written = null;
      }

      emit('ipc-write', {
        api: 'WriteFile',
        handle: this.handle.toString(),
        path: this.path,
        requestedBytes: this.length,
        resultBytes: written,
        retval: retval.toInt32(),
        payload: payloadSummary(this.buffer, Math.max(0, written === null ? this.length : written))
      });
    }
  }
});

hookExport('kernel32.dll', 'ReadFile', {
  onEnter(args) {
    this.handle = args[0];
    this.path = handles.get(this.handle.toString());
    this.buffer = args[1];
    this.length = args[2].toInt32();
    this.readPointer = args[3];
  },
  onLeave(retval) {
    if (this.path !== undefined) {
      let read = null;
      try {
        read = this.readPointer.isNull() ? null : this.readPointer.readU32();
      } catch (_) {
        read = null;
      }

      emit('ipc-read', {
        api: 'ReadFile',
        handle: this.handle.toString(),
        path: this.path,
        requestedBytes: this.length,
        resultBytes: read,
        retval: retval.toInt32(),
        payload: payloadSummary(this.buffer, Math.max(0, read === null ? 0 : read))
      });
    }
  }
});

hookExport('ntdll.dll', 'NtCreateFile', {
  onEnter(args) {
    this.fileHandlePointer = args[0];
    this.path = readObjectAttributesName(args[2]);
    this.interesting = isInterestingPath(this.path);
  },
  onLeave(retval) {
    if (this.interesting && retval.toInt32() >= 0) {
      try {
        const handle = this.fileHandlePointer.readPointer();
        handles.set(handle.toString(), this.path);
        emit('ipc-open', {
          api: 'NtCreateFile',
          handle: handle.toString(),
          path: this.path,
          ntstatus: retval.toString()
        });
      } catch (e) {
        emit('ipc-open-error', { api: 'NtCreateFile', path: this.path, error: String(e) });
      }
    }
  }
});

hookExport('ntdll.dll', 'NtWriteFile', {
  onEnter(args) {
    this.handle = args[0];
    this.path = handles.get(this.handle.toString());
    this.buffer = args[5];
    this.length = args[6].toInt32();
  },
  onLeave(retval) {
    if (this.path !== undefined) {
      emit('ipc-write', {
        api: 'NtWriteFile',
        handle: this.handle.toString(),
        path: this.path,
        requestedBytes: this.length,
        ntstatus: retval.toString(),
        payload: payloadSummary(this.buffer, this.length)
      });
    }
  }
});

hookExport('ntdll.dll', 'NtReadFile', {
  onEnter(args) {
    this.handle = args[0];
    this.path = handles.get(this.handle.toString());
    this.buffer = args[5];
    this.length = args[6].toInt32();
  },
  onLeave(retval) {
    if (this.path !== undefined) {
      emit('ipc-read', {
        api: 'NtReadFile',
        handle: this.handle.toString(),
        path: this.path,
        requestedBytes: this.length,
        ntstatus: retval.toString(),
        payload: payloadSummary(this.buffer, retval.toInt32() >= 0 ? this.length : 0)
      });
    }
  }
});