From a7d5177100cfbe665b6bcb592c9db10f62790b8c Mon Sep 17 00:00:00 2001 From: Yuzu Date: Mon, 13 Jul 2026 16:59:58 +0700 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat:=20support=20multiple=20frame?= =?UTF-8?q?=20sets=20and=20harden=20the=20SSH=20server?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/index.ts | 169 ++++++++++++++++++++++---- src/server.test.ts | 24 ++++ src/server.ts | 291 ++++++++++++++++++++++++++++++++++----------- 3 files changed, 390 insertions(+), 94 deletions(-) diff --git a/src/index.ts b/src/index.ts index a08c39a..604bf2a 100644 --- a/src/index.ts +++ b/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 { + const videoPath = resolveVideoPath(videoArg ?? config.videoPath); + if (!videoPath) { + fail( + `No source video found. Pass --video , 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 { + 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 ` + ); + } + + 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)); +}); diff --git a/src/server.test.ts b/src/server.test.ts index 11a2c29..0015a4f 100644 --- a/src/server.test.ts +++ b/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); + }); }); diff --git a/src/server.ts b/src/server.ts index 83aa08e..5d2b5b0 100644 --- a/src/server.ts +++ b/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 = {}; + 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 | 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 | undefined; + let handshakeTimer: ReturnType | 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 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, - width, - height, - keepAspectRatio - ); - const ascii = frameToAscii( - resized, - config.brightnessThreshold - ); + try { + const ascii = await renderer.render( + currentFrame, + width, + height, + keepAspectRatio + ); - stream.write("ok\x1Bc[0G"); - stream.write("\x1b[2J\x1b[0f"); - stream.write(ascii); + if (ended || stream.destroyed) return; + + stream.write("\x1b[2J\x1b[0f"); + stream.write(ascii); + + currentFrame++; + if (currentFrame < videoData.frames.length) { + return; + } - currentFrame++; - if (currentFrame >= videoData.frames.length) { 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); });