#!/usr/bin/env node

/**
 * PSK Proxy Relay-Node (Server)
 *
 * Exposes two ports:
 * 1. Client port: Listens for TLS-PSK tunnel connections from proxy clients
 * 2. Exit port: Listens for TLS-PSK tunnel connections from exit nodes
 * 
 * Relays frames between connected clients and exits.
 */

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

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

const options = program.opts();

let pskKey;
try {
  pskKey = fs.readFileSync(options.pskFile);
} 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);
}

// Client PSK callback
const clientPskCallback = (socket, identity) => {
  console.log(`Client identity: ${identity}`);
  
  if (identity !== options.clientIdentity) {
    console.warn(`Client PSK identity mismatch. Expected '${options.clientIdentity}', got '${identity}'.`);
    return null;
  }
  
  return pskKey;
};

// Exit PSK callback
const exitPskCallback = (socket, identity) => {
  console.log(`Exit identity: ${identity}`);
  
  if (identity !== options.exitIdentity) {
    console.warn(`Exit PSK identity mismatch. Expected '${options.exitIdentity}', got '${identity}'.`);
    return null;
  }
  
  return pskKey;
};

// Create server for client connections
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);
    });
  }
);

// Create server for exit connections
const exitServer = tls.createServer(
  {
    pskCallback: exitPskCallback,
    ciphers: 'PSK-AES256-GCM-SHA384:PSK-AES128-GCM-SHA256',
  },
  (socket) => {
    if (exitSocket) {
      console.log('Rejecting new exit connection, already have one.');
      socket.destroy();
      return;
    }
    console.log('Exit node connected');
    exitSocket = socket;

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

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

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

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

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

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

// Start listening for exit connections
exitServer.listen(parseInt(options.exitPort, 10), options.host, () => {
  console.log(`PSK Proxy Relay-Node listening for exit nodes on ${options.host}:${options.exitPort}`);
});

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