const net = require("net");
const ws = require("ws");
const tls = require("tls");

const PORT = 8081;
const wss = new ws.WebSocketServer({ port: PORT });

const host = process.argv[2] || process.env['HOST'];
const port = process.argv[3] || process.env['PORT'];
if (!host || !port) {
    console.error("Please specify host and port");
    process.exit(1);
}

const log = process.env['NODE_ENV'] === 'development' ? console.log : () => {};

async function handleClient(ws) {
    let serverSocket;
    const upgradedStatus = { tls: false, secureSocket: null };
    try {
        serverSocket = net.connect({ host, port });
    } catch (e) {
        console.error(`Connect to remote failed: ${e}`);
        ws.send(`Connect to remote failed: ${e}`);
        ws.close();
        return;
    }
    try {
        serverSocket.on('error', (err) => {
            console.error(`Connect to remote failed: ${err}`);
            ws.send(`Connect to remote failed: ${err}`);
            ws.close();
        });

        serverSocket.on('data', (data) => {
            log("Plain. recv socket:", data);
            if (!upgradedStatus.tls)
                ws.send(data, { binary: true });
        });

        ws.on('message', (msg, isBinary) => {
            log("Plain. recv ws:", msg);
            if (!isBinary && msg == 'SSL') {
                log("Upgrade to SSL");
                // Upgrade to SSL/TLS
                upgradedStatus.secureSocket = tls.connect({
                    socket: serverSocket,
                    rejectUnauthorized: false
                });
                
                upgradedStatus.secureSocket.on('secureConnect', () => {
                    const certificate = upgradedStatus.secureSocket.getPeerCertificate(true);
                    ws.send(certificate.raw, { binary: true });
                    upgradedStatus.tls = true;
                });

                upgradedStatus.secureSocket.on('data', (data) => {
                    log("Upgraded to SSL. Send ws:", data);
                    if (upgradedStatus.tls)
                        ws.send(data, { binary: true });
                });

                upgradedStatus.secureSocket.on('error', (error) => {
                    log("Upgraded to SSL. error:", error);
                    ws.close();
                });

                ws.on('message', (msg) => {
                    log("Upgraded to SSL. Send socket:", msg);
                    upgradedStatus.secureSocket.write(msg);
                });

            } else {
                if (!upgradedStatus.tls)
                    serverSocket.write(msg);
            }
        });

        ws.on('close', () => {
            console.info('Client close');
            serverSocket.end();
        });

        serverSocket.on('end', () => {
            console.info('Server close');
            ws.close();
        });
    } catch (error) {
        console.error('An error occurred:', error);
        ws.close();
    }
}

wss.on('connection', handleClient);
console.log("Server is running on port " + PORT)