#!/usr/bin/env node

/**
 * PSK Proxy Relay-Node (Server)
 *
 * Listens for a TLS-PSK tunnel connection from the proxy client.
 * Establishes a single TLS-PSK tunnel connection to an exit node.
 * Relays frames between the client and the exit node.
 */

const tls = require('tls');
const fs = require('fs');
const { program } = require('commander');

program
  .requiredOption('-P, --tunnel-port <port>', 'Port for proxy client TLS-PSK tunnel connections')
  .requiredOption('-H, --host <host>', 'Host to bind to (e.g., 0.0.0.0)')
  .requiredOption('--psk-file <path>', 'Path to PSK key file for client and exit connections')
  .requiredOption('--exit-host <host>', 'Exit node host')
  .requiredOption('--exit-port <port>', 'Exit node port')
  .requiredOption('--exit-identity <identity>', 'PSK identity for the exit node')
  .parse();

const options = program.opts();

let pskKey;
try {
  pskKey = fs.readFileSync(options.pskFile, 'utf8').trim();
} catch (error) {
  console.error(`Error reading PSK file: ${error.message}`);
  process.exit(1);
}

// Global state
let clientSocket = null;
let exitSocket = null;
let clientReader = null;
let exitReader = null;

function createMessageReader(onMessage) {
  let buffer = Buffer.alloc(0);
  let expectedLength = 9;
  let currentMessage = null;

  return function onData(data) {
    buffer = Buffer.concat([buffer, data]);

    while (buffer.length >= expectedLength) {
      if (!currentMessage) {
        const type = buffer.readUInt8(0);
        const connectionId = buffer.readUInt32BE(1);
        const dataLength = buffer.readUInt32BE(5);
        currentMessage = { type, connectionId, dataLength };
        expectedLength = 9 + dataLength;

        if (dataLength === 0) {
          onMessage(currentMessage.type, currentMessage.connectionId, Buffer.alloc(0));
          buffer = buffer.subarray(9);
          currentMessage = null;
          expectedLength = 9;
        }
      } else {
        const messageData = buffer.subarray(9, expectedLength);
        onMessage(currentMessage.type, currentMessage.connectionId, messageData);
        buffer = buffer.subarray(expectedLength);
        currentMessage = null;
        expectedLength = 9;
      }
    }
  };
}

function writeMessage(socket, type, connectionId, data = Buffer.alloc(0)) {
  if (!socket || socket.destroyed) return;
  const buf = Buffer.allocUnsafe(9 + data.length);
  buf.writeUInt8(type, 0);
  buf.writeUInt32BE(connectionId >>> 0, 1);
  buf.writeUInt32BE(data.length >>> 0, 5);
  if (data.length > 0) data.copy(buf, 9);
  socket.write(buf);
}

function connectToExitNode() {
  console.log(`Connecting to exit node ${options.exitHost}:${options.exitPort}...`);

  const pskCb = () => ({
    identity: options.exitIdentity,
    psk: Buffer.from(pskKey, 'hex'),
  });

  const sock = tls.connect(
    {
      host: options.exitHost,
      port: parseInt(options.exitPort, 10),
      pskCallback: pskCb,
      ciphers: 'PSK-AES256-GCM-SHA384:PSK-AES128-GCM-SHA256',
      checkServerIdentity: () => undefined,
    },
    () => {
      console.log('Connected to exit node');
    }
  );

  sock.setNoDelay(true);
  sock.setKeepAlive(true, 30000);

  exitSocket = sock;
  exitReader = createMessageReader((type, connId, data) => {
    // Forward messages from exit to client
    if (clientSocket && !clientSocket.destroyed) {
      writeMessage(clientSocket, type, connId, data);
    }
  });

  sock.on('data', (data) => exitReader(data));

  sock.on('close', () => {
    console.log('Disconnected from exit node. Retrying in 2s...');
    exitSocket = null;
    // If client is still here, it will be disconnected by the client server logic
    if (clientSocket) {
        try { clientSocket.destroy(); } catch (_) {}
    }
    setTimeout(connectToExitNode, 2000);
  });

  sock.on('error', (err) => {
    console.error('Exit node connection error:', err.message);
  });
}

const clientPskCallback = (socket, identity) => {
  console.log(`Client identity: ${identity}`);
  return Buffer.from(pskKey, 'hex');
};

const clientServer = tls.createServer(
  {
    pskCallback: clientPskCallback,
    ciphers: 'PSK-AES256-GCM-SHA384:PSK-AES128-GCM-SHA256',
  },
  (socket) => {
    if (clientSocket) {
      console.log('Rejecting new client connection, already have one.');
      socket.destroy();
      return;
    }
    console.log('Client connected');
    clientSocket = socket;

    socket.setNoDelay(true);
    socket.setKeepAlive(true, 30000);

    clientReader = createMessageReader((type, connId, data) => {
      // Forward messages from client to exit
      if (exitSocket && !exitSocket.destroyed) {
        writeMessage(exitSocket, type, connId, data);
      }
    });

    socket.on('data', (data) => clientReader(data));

    socket.on('close', () => {
      console.log('Client disconnected');
      clientSocket = null;
      // When client disconnects, we can optionally close the exit connection
      // or keep it alive. For simplicity, we keep it.
    });

    socket.on('error', (err) => {
      console.error('Client socket error:', err.message);
    });
  }
);

clientServer.listen(parseInt(options.tunnelPort, 10), options.host, () => {
  console.log(`PSK Proxy Relay-Node listening for clients on ${options.host}:${options.tunnelPort}`);
});

// Start connection to the exit node
connectToExitNode();

process.on('SIGINT', () => {
  console.log('Shutting down...');
  try { clientServer.close(); } catch (_) {}
  if (clientSocket) try { clientSocket.destroy(); } catch (_) {}
  if (exitSocket) try { exitSocket.destroy(); } catch (_) {}
  process.exit(0);
});