app.js

import express from 'express';
import { spawn, execSync } from 'child_process';
import crypto from 'crypto';
import fs from 'fs';
import os from 'os';
import path from 'path';
import { fileURLToPath } from 'url';

const __dirname = path.dirname(fileURLToPath(import.meta.url));

const app = express();
const PORT = process.env.PORT || 7860;
const SECRET_KEY = process.env.SECRET_KEY

if (!SECRET_KEY) {
  console.error('[ERROR] SECRET_KEY environment variable is not set. Exiting.');
  process.exit(1);
}

// SSH password: deterministic from SECRET_KEY — same across restarts, easy to copy
const SSH_PASSWORD = crypto
  .createHash('sha256')
  .update(`ssh:${SECRET_KEY}`)
  .digest('hex')
  .slice(0, 20);

app.use(express.json());
app.use(express.static(path.join(__dirname, 'public')));

// ── Session store ─────────────────────────────────────────────────────────────
const sessions = new Set();

const AUTO_TOKEN = crypto.createHash('sha256').update(`auto:${SECRET_KEY}`).digest('hex');
sessions.add(AUTO_TOKEN);

app.get('/api/auto-token', (_req, res) => {
  res.json({ token: AUTO_TOKEN });
});

// ── ANSI stripper ─────────────────────────────────────────────────────────────
function stripAnsi(str) {
  // eslint-disable-next-line no-control-regex
  return str.replace(/[\u001b\u009b][[()#;?]*(?:[0-9]{1,4}(?:;[0-9]{0,4})*)?[0-9A-ORZcf-nqry=><~]/g, '');
}

// ── Pre-seed shellular config from env vars ───────────────────────────────────
function seedShellularConfig() {
  const hostId    = process.env.SHELLULAR_HOST_ID;
  const keyB64    = process.env.SHELLULAR_KEY;
  const machineId = process.env.SHELLULAR_MACHINE_ID;

  if (!hostId || !keyB64 || !machineId) return;

  const shellularDir = path.join(os.homedir(), '.shellular');
  const configFile   = path.join(shellularDir, 'config.json');
  const keyFile      = path.join(shellularDir, `shellular-${machineId}.e2ee`);

  try {
    fs.mkdirSync(shellularDir, { recursive: true });

    if (!fs.existsSync(configFile)) {
      fs.writeFileSync(configFile, JSON.stringify({ hostId, machineId }), 'utf-8');
      console.log(`[shellular] seeded config: hostId=${hostId}`);
    }

    if (!fs.existsSync(keyFile)) {
      fs.writeFileSync(keyFile, Buffer.from(keyB64, 'base64'), { mode: 0o600 });
      console.log(`[shellular] seeded key: ${keyFile}`);
    }
  } catch (err) {
    console.error('[shellular] failed to seed config:', err.message);
  }
}

seedShellularConfig();

// ── Shellular machine-id helper ───────────────────────────────────────────────
function getHashedMachineId() {
  try {
    const raw = fs.readFileSync('/etc/machine-id', 'utf-8').trim();
    return crypto.createHash('sha256').update(raw).digest('hex');
  } catch {
    return null;
  }
}

app.get('/api/shellular/machine-id', (_req, res) => {
  const id = getHashedMachineId();
  id ? res.json({ machineId: id }) : res.status(500).json({ error: 'Cannot read machine-id' });
});

app.post('/api/shellular/seed-host', requireAuth, (req, res) => {
  const { hostId } = req.body || {};
  if (!hostId || typeof hostId !== 'string' || !hostId.trim()) {
    return res.status(400).json({ error: 'hostId is required' });
  }
  const machineId = getHashedMachineId();
  if (!machineId) return res.status(500).json({ error: 'Cannot read machine-id' });

  try {
    const shellularDir = path.join(os.homedir(), '.shellular');
    fs.mkdirSync(shellularDir, { recursive: true });
    fs.writeFileSync(
      path.join(shellularDir, 'config.json'),
      JSON.stringify({ hostId: hostId.trim(), machineId }, null, 2),
      'utf-8'
    );

    stopShellular();
    outputBuffer = '';
    broadcast({ type: 'clear' });
    setTimeout(startShellular, 600);

    res.json({ ok: true });
  } catch (err) {
    res.status(500).json({ error: err.message });
  }
});

// ── Auth routes ───────────────────────────────────────────────────────────────
app.post('/api/login', (req, res) => {
  const { key } = req.body;
  if (!key || key !== SECRET_KEY) {
    return res.status(401).json({ error: 'Invalid key' });
  }
  const token = crypto.randomUUID();
  sessions.add(token);
  res.json({ token });
});

app.post('/api/logout', (req, res) => {
  const token = req.headers.authorization?.split(' ')[1];
  sessions.delete(token);
  res.json({ ok: true });
});

// ── Auth middleware ────────────────────────────────────────────────────────────
function requireAuth(req, res, next) {
  const token = req.headers.authorization?.split(' ')[1];
  if (!token || !sessions.has(token)) {
    return res.status(401).json({ error: 'Unauthorized' });
  }
  next();
}

// ── Shellular process management ──────────────────────────────────────────────
let shellularProc = null;
let outputBuffer  = '';
const sseClients  = new Set();

function send(res, payload) {
  res.write(`data: ${JSON.stringify(payload)}\n\n`);
}

function broadcast(payload) {
  const frame = `data: ${JSON.stringify(payload)}\n\n`;
  for (const client of sseClients) {
    client.write(frame);
  }
}

let retryTimer = null;

function startShellular() {
  if (shellularProc || retryTimer) return;

  broadcast({ type: 'status', status: 'starting' });

  shellularProc = spawn('shellular', ['--unknown-clients', 'always-allow'], {
    env: { ...process.env, FORCE_COLOR: '0' },
    stdio: ['ignore', 'pipe', 'pipe'],
  });

  let procOutput = '';

  const handleData = (chunk) => {
    const text = stripAnsi(chunk.toString());
    procOutput   += text;
    outputBuffer += text;
    broadcast({ type: 'output', text });
  };

  shellularProc.stdout.on('data', handleData);
  shellularProc.stderr.on('data', handleData);

  shellularProc.on('error', (err) => {
    const text = `\n[spawn error] ${err.message}\n`;
    outputBuffer += text;
    broadcast({ type: 'output', text });
    shellularProc = null;
    broadcast({ type: 'status', status: 'error' });
  });

  shellularProc.on('exit', (code, signal) => {
    shellularProc = null;

    const isRegError = code === 1 && !signal &&
      (procOutput.includes('invalid_union') || procOutput.includes('Too many requests') ||
       procOutput.includes('host registration'));

    if (isRegError) {
      const WAIT = 30;
      const msg = `\n⚠ Registration rate-limited by shellular API.\n` +
                  `  Retrying automatically in ${WAIT}s — please wait…\n`;
      outputBuffer += msg;
      broadcast({ type: 'output', text: msg });
      broadcast({ type: 'status', status: 'retrying' });

      retryTimer = setTimeout(() => {
        retryTimer = null;
        const msg2 = '\n[Retrying registration…]\n';
        outputBuffer += msg2;
        broadcast({ type: 'output', text: msg2 });
        startShellular();
      }, WAIT * 1000);
    } else {
      const text = code !== 0
        ? `\n[shellular exited — code=${code ?? '?'}, signal=${signal ?? 'none'}]\n`
        : '\n[shellular disconnected]\n';
      outputBuffer += text;
      broadcast({ type: 'output', text });
      broadcast({ type: 'status', status: 'stopped' });
    }
  });

  broadcast({ type: 'status', status: 'running' });
}

function stopShellular() {
  if (retryTimer) { clearTimeout(retryTimer); retryTimer = null; }
  if (!shellularProc) return;
  shellularProc.kill('SIGTERM');
  shellularProc = null;
}

// ── Python sync.py subprocess ─────────────────────────────────────────────────
let syncProc = null;

function startSyncPy() {
  if (syncProc) return;

  console.log('[sync] Starting sync.py...');

  syncProc = spawn('python3', [path.join(__dirname, 'syn.py')], {
    env: { ...process.env },
    stdio: ['ignore', 'pipe', 'pipe'],
  });

  syncProc.stdout.on('data', (chunk) => {
    console.log('[sync]', chunk.toString().trim());
  });

  syncProc.stderr.on('data', (chunk) => {
    console.error('[sync:err]', chunk.toString().trim());
  });

  syncProc.on('error', (err) => {
    console.error('[sync] Spawn error:', err.message);
    syncProc = null;
  });

  syncProc.on('exit', (code, signal) => {
    console.warn(`[sync] sync.py exited — code=${code ?? '?'}, signal=${signal ?? 'none'}`);
    syncProc = null;
    setTimeout(startSyncPy, 10_000);
  });
}

function stopSyncPy() {
  if (!syncProc) return;
  syncProc.kill('SIGTERM');
  syncProc = null;
}

// ── SSH server setup ──────────────────────────────────────────────────────────
function setupSSH() {
  try {
    // Set root password
    execSync(`echo "root:${SSH_PASSWORD}" | chpasswd`, { stdio: 'ignore' });
    console.log('[ssh] root password set');

    // Allow root + password login
    fs.mkdirSync('/etc/ssh/sshd_config.d', { recursive: true });
    fs.writeFileSync('/etc/ssh/sshd_config.d/99-termius.conf', [
      'PermitRootLogin yes',
      'PasswordAuthentication yes',
      'ChallengeResponseAuthentication no',
      'UsePAM no',
      'PrintMotd no',
    ].join('\n') + '\n', 'utf-8');

    // Generate host keys (no-op if already present)
    execSync('ssh-keygen -A', { stdio: 'ignore' });

    // Launch sshd
    const sshd = spawn('/usr/sbin/sshd', ['-D', '-e'], {
      detached: true,
      stdio: 'ignore',
    });
    sshd.unref();
    console.log('[ssh] sshd started, pid:', sshd.pid);

    // Give sshd a moment then open the bore tunnel
    setTimeout(startBore, 2000);
  } catch (err) {
    console.error('[ssh] setup error:', err.message);
  }
}

// ── Bore tunnel (exposes SSH port to the internet) ────────────────────────────
let boreProc = null;
let boreHost = null;
let borePort = null;

function startBore() {
  if (boreProc) return;
  console.log('[bore] starting tunnel…');

  boreProc = spawn('bore', ['local', '22', '--to', 'bore.pub'], {
    stdio: ['ignore', 'pipe', 'pipe'],
  });

  const onData = (chunk) => {
    const text = chunk.toString();
    console.log('[bore]', text.trim());
    // bore prints: "… Listening at bore.pub:NNNNN"
    const m = text.match(/bore\.pub:(\d+)/i);
    if (m) {
      boreHost = 'bore.pub';
      borePort = parseInt(m[1], 10);
      console.log(`[bore] tunnel ready → ${boreHost}:${borePort}`);
    }
  };

  boreProc.stdout.on('data', onData);
  boreProc.stderr.on('data', onData);

  boreProc.on('error', (err) => {
    console.error('[bore] error:', err.message);
    boreProc = null; boreHost = null; borePort = null;
    setTimeout(startBore, 15_000);
  });

  boreProc.on('exit', (code, signal) => {
    console.warn(`[bore] exited code=${code} signal=${signal} — restarting in 10 s`);
    boreProc = null; boreHost = null; borePort = null;
    setTimeout(startBore, 10_000);
  });
}

function stopBore() {
  if (!boreProc) return;
  boreProc.kill('SIGTERM');
  boreProc = null; boreHost = null; borePort = null;
}

// ── SSH info endpoint (used by frontend to show Termius credentials) ──────────
app.get('/api/ssh-info', requireAuth, (_req, res) => {
  res.json({
    ready:    !!(boreHost && borePort),
    host:     boreHost,
    port:     borePort,
    username: 'root',
    password: SSH_PASSWORD,
  });
});

// ── SSE stream ─────────────────────────────────────────────────────────────────
app.get('/api/stream', requireAuth, (req, res) => {
  res.setHeader('Content-Type', 'text/event-stream');
  res.setHeader('Cache-Control', 'no-cache');
  res.setHeader('Connection', 'keep-alive');
  res.setHeader('X-Accel-Buffering', 'no');
  res.flushHeaders();

  send(res, { type: 'status', status: shellularProc ? 'running' : 'stopped' });

  if (outputBuffer) {
    send(res, { type: 'output', text: outputBuffer });
  }

  sseClients.add(res);
  req.on('close', () => sseClients.delete(res));
});

// ── Control endpoints ──────────────────────────────────────────────────────────
app.post('/api/shellular/start', requireAuth, (_req, res) => {
  startShellular();
  res.json({ ok: true, running: !!shellularProc });
});

app.post('/api/shellular/stop', requireAuth, (_req, res) => {
  stopShellular();
  outputBuffer = '';
  broadcast({ type: 'output', text: '' });
  res.json({ ok: true });
});

app.post('/api/shellular/restart', requireAuth, (_req, res) => {
  stopShellular();
  outputBuffer = '';
  broadcast({ type: 'clear' });
  setTimeout(startShellular, 600);
  res.json({ ok: true });
});

app.get('/api/status', requireAuth, (_req, res) => {
  res.json({ running: !!shellularProc });
});

app.get('/api/setup-status', requireAuth, (_req, res) => {
  const seeded = !!(
    process.env.SHELLULAR_HOST_ID &&
    process.env.SHELLULAR_KEY &&
    process.env.SHELLULAR_MACHINE_ID
  );
  res.json({ seeded });
});

app.get('/api/shellular/credentials', requireAuth, (_req, res) => {
  try {
    const shellularDir = path.join(os.homedir(), '.shellular');
    const configRaw    = fs.readFileSync(path.join(shellularDir, 'config.json'), 'utf-8');
    const { hostId, machineId } = JSON.parse(configRaw);
    const keyFile = path.join(shellularDir, `shellular-${machineId}.e2ee`);
    const keyB64  = fs.readFileSync(keyFile).toString('base64');
    res.json({ hostId, machineId, keyB64 });
  } catch {
    res.status(404).json({ error: 'Not registered yet.' });
  }
});

app.get('/api/shellular/qr-data', requireAuth, (_req, res) => {
  try {
    const shellularDir = path.join(os.homedir(), '.shellular');
    const configRaw    = fs.readFileSync(path.join(shellularDir, 'config.json'), 'utf-8');
    const { hostId, machineId } = JSON.parse(configRaw);
    const keyFile = path.join(shellularDir, `shellular-${machineId}.e2ee`);
    const keyB64  = fs.readFileSync(keyFile).toString('base64');
    res.json({ qrData: `${hostId}:${keyB64}` });
  } catch {
    res.status(404).json({ error: 'Config not seeded yet.' });
  }
});

// ── Start ──────────────────────────────────────────────────────────────────────
app.listen(PORT, '0.0.0.0', () => {
  console.log(`Shellular Web UI → http://0.0.0.0:${PORT}`);
  startSyncPy();    // 🐍 Sync to HF dataset
  startShellular(); // 🚀 Shellular for QR access
  setupSSH();       // 🔒 SSH server + bore tunnel for Termius
});

// ── Graceful shutdown ──────────────────────────────────────────────────────────
process.on('SIGTERM', () => { stopShellular(); stopSyncPy(); stopBore(); });
process.on('SIGINT',  () => { stopShellular(); stopSyncPy(); stopBore(); });