✨ feat: support multiple frame sets and harden the SSH server

src/index.ts

@@ -1,49 +1,172 @@
 import fs from "fs";
+import os from "os";
 import path from "path";
-import { loadConfig, loadOptionalTextFile } from "./config";
+import { loadConfig, loadOptionalTextFile, Config } from "./config";
 import { ensureHostKeys } from "./hostKeys";
-import { loadFrames } from "./frames";
+import { loadFramesAsync, FramesContainer } from "./frames";
 import { createServer } from "./server";
 import videoProcessor from "./videoProcessor";
+import { logger } from "./logger";
+
+const DATA_DIR = path.join(process.cwd(), "data");
+const FRAMES_DIR = path.join(process.cwd(), "frames");
+
+function fail(message: string): never {
+    logger.error(message);
+    process.exit(1);
+}
+
+interface Args {
+    generate: boolean;
+    video?: string;
+}
+
+function parseArgs(argv: string[]): Args {
+    const args: Args = { generate: false };
+    for (let i = 0; i < argv.length; i++) {
+        const arg = argv[i];
+        if (arg === "--generate" || arg === "-g") args.generate = true;
+        else if (arg === "--video" || arg === "-v") args.video = argv[++i];
+    }
+    return args;
+}
+
+function resolveVideoPath(explicitPath?: string): string | undefined {
+    if (explicitPath) {
+        const explicit = path.resolve(explicitPath);
+        return fs.existsSync(explicit) ? explicit : undefined;
+    }
+
+    const cwd = process.cwd();
+    const match = fs
+        .readdirSync(cwd)
+        .filter((name) => path.parse(name).name.toLowerCase() === "video")
+        .sort()
+        .find((name) => fs.statSync(path.join(cwd, name)).isFile());
+
+    return match ? path.join(cwd, match) : undefined;
+}
+
+async function generateFrames(
+    config: Config,
+    videoArg?: string
+): Promise<void> {
+    const videoPath = resolveVideoPath(videoArg ?? config.videoPath);
+    if (!videoPath) {
+        fail(
+            `No source video found. Pass --video <path>, set VIDEO_PATH, or ` +
+                `drop a "video.*" file in "${process.cwd()}".`
+        );
+    }
+
+    fs.mkdirSync(FRAMES_DIR, { recursive: true });
+    const output = path.join(FRAMES_DIR, `${path.parse(videoPath).name}.json`);
+
+    logger.info(`Generating frames from "${videoPath}" -> ${output}`);
+    try {
+        await videoProcessor.process(videoPath, output, {
+            maxDimension: config.frameResolution,
+        });
+    } catch (err) {
+        fail(
+            `Failed to generate frames from "${videoPath}": ` +
+                (err instanceof Error ? err.message : String(err))
+        );
+    }
+}
+
+async function loadAllFrames(): Promise<FramesContainer[]> {
+    const files = fs.existsSync(FRAMES_DIR)
+        ? fs
+              .readdirSync(FRAMES_DIR)
+              .filter((name) => name.toLowerCase().endsWith(".json"))
+              .sort()
+        : [];
+
+    if (files.length === 0) {
+        fail(
+            `No frame sets found in "${FRAMES_DIR}". ` +
+                `Generate one first with: bun src/index.ts --generate --video <path>`
+        );
+    }
+
+    const cpus = os.availableParallelism?.() ?? os.cpus().length;
+    const concurrency = Math.min(files.length, Math.max(1, Math.min(cpus, 4)));
+
+    const results: FramesContainer[] = new Array(files.length);
+    let nextIndex = 0;
+    const worker = async () => {
+        for (let i = nextIndex++; i < files.length; i = nextIndex++) {
+            const file = files[i];
+            const filePath = path.join(FRAMES_DIR, file);
+            const sizeMb = (fs.statSync(filePath).size / 1024 / 1024).toFixed(
+                1
+            );
+            logger.info(`Loading ${file} (${sizeMb} MB)...`);
+            const data = await loadFramesAsync(filePath);
+            data.name = file;
+            logger.info(
+                `  ${file}: ${data.frames.length} frames @ ${data.fps}fps`
+            );
+            results[i] = data;
+        }
+    };
+
+    await Promise.all(Array.from({ length: concurrency }, () => worker()));
+    return results;
+}
 
 async function main() {
     const config = loadConfig();
-    const configDir = path.join(process.cwd(), "config");
+    const args = parseArgs(process.argv.slice(2));
 
-    if (!fs.existsSync(configDir)) {
-        fs.mkdirSync(configDir);
+    if (args.generate) {
+        await generateFrames(config, args.video);
+        return;
     }
 
-    const bannerText = loadOptionalTextFile(path.join(configDir, "banner.txt"));
+    fs.mkdirSync(DATA_DIR, { recursive: true });
+
+    const bannerText = loadOptionalTextFile(path.join(DATA_DIR, "banner.txt"));
     const fakeLoginText = loadOptionalTextFile(
-        path.join(configDir, "fakelogin.txt")
+        path.join(DATA_DIR, "fakelogin.txt")
     );
     const goodbyeText = loadOptionalTextFile(
-        path.join(configDir, "goodbye.txt")
+        path.join(DATA_DIR, "goodbye.txt")
     );
 
-    ensureHostKeys(configDir);
-    const hostKey = fs.readFileSync(path.join(configDir, "id_rsa"));
-
-    const framesPath = path.join(configDir, "frames.json");
-    if (!fs.existsSync(framesPath)) {
-        console.error("frames.json not found!, generating frames.json...");
-        await videoProcessor.process("video.mp4", framesPath);
-    }
-
-    const videoData = loadFrames(framesPath);
-    console.log("Loaded frames");
+    const hostKeys = ensureHostKeys(DATA_DIR);
+    const videoSets = await loadAllFrames();
+    logger.info(`Loaded ${videoSets.length} frame set(s)`);
 
     const server = createServer({
         config,
-        hostKey,
+        hostKeys,
         bannerText,
         fakeLoginText,
         goodbyeText,
-        videoData,
+        videoSets,
     });
 
-    server.listen(config.port, config.host);
+    server.on("error", (err: Error) => {
+        logger.error("Server error:", err.message);
+        process.exit(1);
+    });
+
+    server.listen(config.port, config.host, () => {
+        logger.info(`TrollSSH listening on ${config.host}:${config.port}`);
+    });
+
+    const shutdown = (signal: string) => {
+        logger.info(`Received ${signal}, shutting down...`);
+        server.close(() => process.exit(0));
+        // Fail-safe: force exit if connections don't drain promptly.
+        setTimeout(() => process.exit(0), 5000).unref();
+    };
+    process.on("SIGINT", () => shutdown("SIGINT"));
+    process.on("SIGTERM", () => shutdown("SIGTERM"));
 }
 
-main();
+main().catch((err) => {
+    fail(err instanceof Error ? err.message : String(err));
+});

src/server.test.ts

@@ -44,4 +44,28 @@ describe("ConnectionTracker", () => {
         expect(tracker.count("1.1.1.1")).toBe(1);
         expect(tracker.count("2.2.2.2")).toBe(2);
     });
+
+    test("totalCount aggregates connections across all IPs", () => {
+        const tracker = new ConnectionTracker();
+        tracker.increment("1.1.1.1");
+        tracker.increment("2.2.2.2");
+        expect(tracker.totalCount()).toBe(2);
+        tracker.decrement("1.1.1.1");
+        expect(tracker.totalCount()).toBe(1);
+    });
+
+    test("hasReachedTotalLimit reflects the global total", () => {
+        const tracker = new ConnectionTracker();
+        tracker.increment("1.1.1.1");
+        tracker.increment("2.2.2.2");
+        expect(tracker.hasReachedTotalLimit(2)).toBe(true);
+        expect(tracker.hasReachedTotalLimit(3)).toBe(false);
+    });
+
+    test("decrementing an unknown IP does not affect the total", () => {
+        const tracker = new ConnectionTracker();
+        tracker.increment("1.1.1.1");
+        tracker.decrement("9.9.9.9");
+        expect(tracker.totalCount()).toBe(1);
+    });
 });

src/server.ts

@@ -1,18 +1,24 @@
 import ssh2 from "ssh2";
 import { Config } from "./config";
-import { FramesContainer, resizeFrame, frameToAscii } from "./frames";
+import { FramesContainer, FrameRenderer } from "./frames";
+import { logger, sanitize } from "./logger";
+
+const MAX_WRITE_BACKLOG_BYTES = 4 * 1024 * 1024;
 
 export class ConnectionTracker {
     private counts: Record<string, number> = {};
+    private total = 0;
 
     increment(ip: string): number {
         this.counts[ip] = (this.counts[ip] ?? 0) + 1;
+        this.total += 1;
         return this.counts[ip];
     }
 
     decrement(ip: string): void {
         if (typeof this.counts[ip] === "undefined") return;
         this.counts[ip] -= 1;
+        this.total = Math.max(0, this.total - 1);
         if (this.counts[ip] <= 0) delete this.counts[ip];
     }
 
@@ -20,85 +26,153 @@ export class ConnectionTracker {
         return this.counts[ip] ?? 0;
     }
 
+    totalCount(): number {
+        return this.total;
+    }
+
     hasReachedLimit(ip: string, max: number): boolean {
         return this.count(ip) >= max;
     }
+
+    hasReachedTotalLimit(max: number): boolean {
+        return this.total >= max;
+    }
 }
 
 export interface ServerDeps {
     config: Config;
-    hostKey: Buffer;
+    hostKeys: Buffer[];
     bannerText?: string;
     fakeLoginText?: string;
     goodbyeText?: string;
-    videoData: FramesContainer;
+    videoSets: FramesContainer[];
+}
+
+function clampDimension(value: number, max: number): number {
+    if (!Number.isFinite(value) || value < 1) return 1;
+    return Math.min(Math.floor(value), max);
 }
 
 export function createServer(deps: ServerDeps): ssh2.Server {
     const {
         config,
-        hostKey,
+        hostKeys,
         bannerText,
         fakeLoginText,
         goodbyeText,
-        videoData,
+        videoSets,
     } = deps;
     const tracker = new ConnectionTracker();
 
+    const sets = videoSets.map((data) => ({
+        data,
+        renderer: new FrameRenderer(data.frames, {
+            brightnessThreshold: config.brightnessThreshold,
+            charset: config.charset,
+            invert: config.invert,
+        }),
+    }));
+
     const server = new ssh2.Server({
-        hostKeys: [hostKey],
+        hostKeys,
         banner: bannerText,
     });
 
     server.on("connection", (client, info) => {
-        if (tracker.hasReachedLimit(info.ip, config.maxConnections)) {
+        if (
+            tracker.hasReachedTotalLimit(config.maxTotalConnections) ||
+            tracker.hasReachedLimit(info.ip, config.maxConnections)
+        ) {
             client.on("error", () => {});
             client.end();
+            logger.warn("Connection rejected (limit reached) from", info.ip);
             return;
         }
 
-        tracker.increment(info.ip);
-        console.log("New connection from", info.ip);
+        const activeForIp = tracker.increment(info.ip);
+        let currentSetIndex = Math.floor(Math.random() * sets.length);
+        let { data: videoData, renderer } = sets[currentSetIndex];
 
-        client.on("handshake", () => {
-            console.log("Handshake from", info.ip);
-        });
-
-        let interval: ReturnType<typeof setInterval> | undefined;
-
-        const endSession = () => {
-            tracker.decrement(info.ip);
-            if (interval) clearInterval(interval);
+        const pickNextSetIndex = (exclude: number): number => {
+            if (sets.length <= 1) return exclude;
+            let next = exclude;
+            while (next === exclude) {
+                next = Math.floor(Math.random() * sets.length);
+            }
+            return next;
         };
 
+        logger.info(
+            `New connection from ${info.ip} ` +
+                `(ip=${activeForIp}, total=${tracker.totalCount()}) ` +
+                `-> playing "${videoData.name ?? "?"}"`
+        );
+
+        let interval: ReturnType<typeof setInterval> | undefined;
+        let handshakeTimer: ReturnType<typeof setTimeout> | undefined;
+        let authAttempts = 0;
+        let ended = false;
+
+        // Force-drop clients that connect but never complete a handshake
+        if (config.handshakeTimeout > 0) {
+            handshakeTimer = setTimeout(() => {
+                logger.warn("Handshake timeout for", info.ip);
+                client.end();
+            }, config.handshakeTimeout);
+        }
+
+        const endSession = () => {
+            if (ended) return;
+            ended = true;
+            tracker.decrement(info.ip);
+            if (interval) clearInterval(interval);
+            if (handshakeTimer) clearTimeout(handshakeTimer);
+        };
+
+        client.on("handshake", () => {
+            logger.debug("Handshake from", info.ip);
+            if (handshakeTimer) {
+                clearTimeout(handshakeTimer);
+                handshakeTimer = undefined;
+            }
+        });
+
         client.on("close", () => {
-            console.log("Client closed connection from", info.ip);
+            logger.info("Client closed connection from", info.ip);
             endSession();
         });
 
         client.on("error", (err) => {
             if (err.message === "read ECONNRESET") {
-                console.log(
+                logger.debug(
                     "Terminal closed (ECONNRESET) for session",
                     info.ip
                 );
-                endSession();
-                return;
+            } else {
+                logger.warn(
+                    `Client error from ${info.ip}:`,
+                    sanitize(err.message)
+                );
             }
-            console.log("Client error: ", err);
+            endSession();
         });
 
         client.on("authentication", (ctx) => {
             if (ctx.method === "password" && config.logCredentials)
-                console.log(
-                    "Authentication from",
-                    info.ip,
-                    ctx.method,
-                    ctx.username,
-                    ctx.password
+                logger.info(
+                    `Auth attempt from ${info.ip} method=${ctx.method} ` +
+                        `user="${sanitize(ctx.username, 128)}" ` +
+                        `pass="${sanitize(ctx.password, 128)}"`
                 );
-            if (!ctx.username) return ctx.reject(["password"]);
+
+            authAttempts += 1;
+            if (authAttempts > config.maxAuthAttempts) {
+                client.end();
+                return;
+            }
+
             if (ctx.method !== "password") return ctx.reject(["password"]);
+            if (!ctx.username) return ctx.reject(["password"]);
             if (!ctx.password) return ctx.reject(["password"]);
 
             ctx.accept();
@@ -107,19 +181,19 @@ export function createServer(deps: ServerDeps): ssh2.Server {
         client.on("session", (accept, _reject) => {
             const session = accept();
 
-            let height = 100;
-            let width = 100;
+            let height = clampDimension(24, config.maxDimension);
+            let width = clampDimension(80, config.maxDimension);
 
             session.once("pty", (accept, _reject, data) => {
-                console.log("Opening pty for session", info.ip);
-                height = data.rows;
-                width = data.cols;
+                logger.debug("Opening pty for session", info.ip);
+                height = clampDimension(data.rows, config.maxDimension);
+                width = clampDimension(data.cols, config.maxDimension);
                 accept();
             });
 
             session.on("window-change", (_accept, _reject, data) => {
-                height = data.rows;
-                width = data.cols;
+                height = clampDimension(data.rows, config.maxDimension);
+                width = clampDimension(data.cols, config.maxDimension);
             });
 
             const playVideo = (
@@ -127,79 +201,154 @@ export function createServer(deps: ServerDeps): ssh2.Server {
                 keepAspectRatio: boolean
             ) => {
                 stream.setEncoding("utf8");
-                console.log("Terminal size: " + width + "x" + height);
+                logger.debug(`Terminal size ${width}x${height} for ${info.ip}`);
 
                 if (typeof fakeLoginText !== "undefined") {
                     stream.write("\x1b[2J\x1b[0f");
                     stream.write(fakeLoginText);
                 }
 
-                setTimeout(() => {
                 let currentFrame = 0;
                 let loopCount = 0;
+                let rendering = false;
 
+                const startRenderLoop = () => {
                     interval = setInterval(async () => {
-                        if (stream.destroyed) {
-                            clearInterval(interval);
+                        if (ended || stream.destroyed) {
+                            if (interval) clearInterval(interval);
                             return;
                         }
+                        if (rendering) return;
+                        if (
+                            (stream.writableLength ?? 0) >
+                            MAX_WRITE_BACKLOG_BYTES
+                        ) {
+                            return;
+                        }
+                        rendering = true;
 
-                        const frame = Buffer.from(
-                            videoData.frames[currentFrame]
-                        );
-                        const resized = await resizeFrame(
-                            frame,
+                        try {
+                            const ascii = await renderer.render(
+                                currentFrame,
                                 width,
                                 height,
                                 keepAspectRatio
                             );
-                        const ascii = frameToAscii(
-                            resized,
-                            config.brightnessThreshold
-                        );
 
-                        stream.write("ok\x1Bc[0G");
+                            if (ended || stream.destroyed) return;
+
                             stream.write("\x1b[2J\x1b[0f");
                             stream.write(ascii);
 
                             currentFrame++;
-                        if (currentFrame >= videoData.frames.length) {
+                            if (currentFrame < videoData.frames.length) {
+                                return;
+                            }
+
                             currentFrame = 0;
                             loopCount++;
                             if (loopCount >= config.maxLoop) {
-                                clearInterval(interval);
+                                if (interval) clearInterval(interval);
                                 stream.write("\x1b[2J\x1b[0f");
                                 if (typeof goodbyeText !== "undefined") {
                                     stream.write(goodbyeText);
                                 }
                                 setTimeout(() => {
-                                    console.log(
-                                        "Terminal closed for session",
+                                    logger.info(
+                                        "Playback finished, closing session",
                                         info.ip
                                     );
-                                    tracker.decrement(info.ip);
                                     stream.end();
                                     client.end();
                                 }, 1000);
+                                return;
                             }
+
+                            if (config.playbackMode === "random") {
+                                currentSetIndex =
+                                    pickNextSetIndex(currentSetIndex);
+                                ({ data: videoData, renderer } =
+                                    sets[currentSetIndex]);
+                                logger.info(
+                                    `Playthrough done for ${info.ip}, ` +
+                                        `switching to "${videoData.name ?? "?"}"`
+                                );
+                                if (interval) clearInterval(interval);
+                                rendering = false;
+                                startRenderLoop();
+                            } else {
+                                logger.info(
+                                    `Playthrough done for ${info.ip}, ` +
+                                        `looping "${videoData.name ?? "?"}" ` +
+                                        `(${loopCount}/${config.maxLoop})`
+                                );
+                            }
+                        } catch (err) {
+                            logger.error(
+                                "Render error for",
+                                info.ip,
+                                sanitize(
+                                    err instanceof Error ? err.message : err
+                                )
+                            );
+                            if (interval) clearInterval(interval);
+                            client.end();
+                        } finally {
+                            rendering = false;
                         }
                     }, 1000 / videoData.fps);
+                };
+
+                let lastSwitch = 0;
+                const switchSet = (delta: number) => {
+                    if (sets.length <= 1) return;
+                    currentSetIndex =
+                        (currentSetIndex + delta + sets.length) % sets.length;
+                    ({ data: videoData, renderer } = sets[currentSetIndex]);
+                    currentFrame = 0;
+                    logger.debug(
+                        `${info.ip} switched to "${videoData.name ?? "?"}"`
+                    );
+                    if (interval) {
+                        clearInterval(interval);
+                        rendering = false;
+                        startRenderLoop();
+                    }
+                };
+
+                if (config.allowUserControl) {
+                    stream.on("data", (chunk: Buffer | string) => {
+                        const s = chunk.toString();
+                        let delta = 0;
+                        if (s.includes("\x1b[C") || s.includes("\x1b[A"))
+                            delta = 1;
+                        else if (s.includes("\x1b[D") || s.includes("\x1b[B"))
+                            delta = -1;
+                        if (delta === 0) return;
+                        const now = Date.now();
+                        if (now - lastSwitch < config.switchDebounceMs) return;
+                        lastSwitch = now;
+                        switchSet(delta);
+                    });
+                }
+
+                setTimeout(() => {
+                    if (ended || stream.destroyed) return;
+                    startRenderLoop();
                 }, config.loginDelay);
             };
 
             session.once("exec", (accept, _reject, data) => {
-                console.log(
-                    "Client",
-                    info.ip,
-                    "is trying to execute command",
-                    '"' + data.command + '"'
+                logger.info(
+                    `Client ${info.ip} attempted exec: ` +
+                        `"${sanitize(data.command, 512)}"`
                 );
                 const stream = accept();
                 playVideo(stream, false);
             });
 
             session.once("shell", (accept, _reject) => {
-                console.log("Opening shell for session", info.ip);
+                logger.debug("Opening shell for session", info.ip);
                 const stream = accept();
                 playVideo(stream, false);
             });