278 lines
9.2 KiB
TypeScript
278 lines
9.2 KiB
TypeScript
import { Bot, Context } from "grammy";
|
|
import Anthropic from "@anthropic-ai/sdk";
|
|
import { readFileSync } from "fs";
|
|
import { resolve } from "path";
|
|
|
|
// ─── Config ───
|
|
const BOT_TOKEN = process.env.TELEGRAM_BOT_TOKEN!;
|
|
const ANTHROPIC_KEY = process.env.ANTHROPIC_API_KEY!;
|
|
const SSH_USER = process.env.SSH_USER || "root";
|
|
const SSH_HOST = process.env.SSH_HOST!;
|
|
const SSH_PORT = process.env.SSH_PORT || "22";
|
|
const ALLOWED_USERS = process.env.TELEGRAM_ALLOWED_USERS?.split(",").map(Number) || [];
|
|
|
|
if (!BOT_TOKEN) throw new Error("TELEGRAM_BOT_TOKEN missing from .env");
|
|
if (!ANTHROPIC_KEY) throw new Error("ANTHROPIC_API_KEY missing from .env");
|
|
if (!SSH_HOST) throw new Error("SSH_HOST missing from .env");
|
|
|
|
// ─── Load infra context ───
|
|
let infraContext = "";
|
|
try {
|
|
infraContext = readFileSync(resolve(__dirname, "../.pi/infra.md"), "utf-8");
|
|
} catch {
|
|
console.warn("⚠️ Could not load .pi/infra.md");
|
|
}
|
|
|
|
// ─── Init ───
|
|
const bot = new Bot(BOT_TOKEN);
|
|
const anthropic = new Anthropic({ apiKey: ANTHROPIC_KEY });
|
|
|
|
// Per-user conversation history
|
|
const conversations = new Map<number, Anthropic.MessageParam[]>();
|
|
const MAX_HISTORY = 30;
|
|
|
|
// ─── Auth middleware ───
|
|
function isAuthorized(ctx: Context): boolean {
|
|
// Reject if no whitelist configured (fail-closed)
|
|
if (ALLOWED_USERS.length === 0) return false;
|
|
return ALLOWED_USERS.includes(ctx.from?.id || 0);
|
|
}
|
|
|
|
// ─── Rate limiting ───
|
|
const rateLimits = new Map<number, number[]>();
|
|
const RATE_LIMIT_WINDOW = 60_000; // 1 minute
|
|
const RATE_LIMIT_MAX = 15; // max messages per window
|
|
|
|
function isRateLimited(userId: number): boolean {
|
|
const now = Date.now();
|
|
const timestamps = (rateLimits.get(userId) || []).filter(t => now - t < RATE_LIMIT_WINDOW);
|
|
timestamps.push(now);
|
|
rateLimits.set(userId, timestamps);
|
|
return timestamps.length > RATE_LIMIT_MAX;
|
|
}
|
|
|
|
// ─── Blocked command patterns ───
|
|
const BLOCKED_PATTERNS = [
|
|
/rm\s+-rf\s+\/(?!\w)/i, // rm -rf /
|
|
/mkfs\./i, // format disk
|
|
/dd\s+if=.*of=\/dev/i, // overwrite disk
|
|
/:\(\)\s*\{\s*:\|:/, // fork bomb
|
|
];
|
|
|
|
function isDangerous(cmd: string): string | null {
|
|
for (const pattern of BLOCKED_PATTERNS) {
|
|
if (pattern.test(cmd)) return `⛔ Blocked: matches dangerous pattern \`${pattern.source}\``;
|
|
}
|
|
return null;
|
|
}
|
|
|
|
// ─── SSH exec helper ───
|
|
async function sshExec(command: string, timeoutSec = 30): Promise<string> {
|
|
const proc = Bun.spawn(
|
|
["ssh", "-o", "ConnectTimeout=5", "-o", "StrictHostKeyChecking=accept-new",
|
|
"-p", SSH_PORT, `${SSH_USER}@${SSH_HOST}`, command],
|
|
{ stdout: "pipe", stderr: "pipe" }
|
|
);
|
|
|
|
const timer = setTimeout(() => proc.kill(), timeoutSec * 1000);
|
|
const [stdout, stderr] = await Promise.all([
|
|
new Response(proc.stdout).text(),
|
|
new Response(proc.stderr).text(),
|
|
]);
|
|
clearTimeout(timer);
|
|
await proc.exited;
|
|
|
|
const output = (stdout + stderr).trim();
|
|
return output.length > 3500 ? output.slice(0, 3500) + "\n...(truncated)" : output;
|
|
}
|
|
|
|
// ─── System prompt ───
|
|
const SYSTEM_PROMPT = `You are a DevOps assistant connected to a live server via Telegram.
|
|
You have access to run SSH commands on the server. You prioritize uptime.
|
|
|
|
Infrastructure context:
|
|
${infraContext}
|
|
|
|
RULES:
|
|
- If the user asks to run a command, use the run_command tool.
|
|
- For destructive commands (rm -rf, DROP, reboot, etc.), WARN the user and ask for confirmation before executing.
|
|
- Keep responses concise — this is Telegram, not a terminal.
|
|
- You can chain multiple commands in one SSH call using && or ;
|
|
- Format output with monospace blocks for readability.`;
|
|
|
|
// ─── Claude tools ───
|
|
const tools: Anthropic.Tool[] = [
|
|
{
|
|
name: "run_command",
|
|
description: "Execute a command on the server via SSH. Use for any server inspection, management, or debugging task.",
|
|
input_schema: {
|
|
type: "object" as const,
|
|
properties: {
|
|
command: { type: "string", description: "Shell command to run on the server" },
|
|
timeout: { type: "number", description: "Timeout in seconds (default 30)" },
|
|
},
|
|
required: ["command"],
|
|
},
|
|
},
|
|
];
|
|
|
|
// ─── Process message through Claude ───
|
|
async function chat(userId: number, message: string): Promise<string> {
|
|
// Get or init conversation history
|
|
if (!conversations.has(userId)) conversations.set(userId, []);
|
|
const history = conversations.get(userId)!;
|
|
|
|
history.push({ role: "user", content: message });
|
|
|
|
// Trim history if too long
|
|
while (history.length > MAX_HISTORY) history.shift();
|
|
|
|
let response = await anthropic.messages.create({
|
|
model: "claude-sonnet-4-20250514",
|
|
max_tokens: 1024,
|
|
system: SYSTEM_PROMPT,
|
|
tools,
|
|
messages: history,
|
|
});
|
|
|
|
// Tool use loop
|
|
while (response.stop_reason === "tool_use") {
|
|
const assistantContent = response.content;
|
|
history.push({ role: "assistant", content: assistantContent });
|
|
|
|
const toolResults: Anthropic.ToolResultBlockParam[] = [];
|
|
|
|
for (const block of assistantContent) {
|
|
if (block.type === "tool_use") {
|
|
const input = block.input as { command: string; timeout?: number };
|
|
console.log(`🔧 [${userId}] run_command: ${input.command}`);
|
|
const danger = isDangerous(input.command);
|
|
if (danger) {
|
|
toolResults.push({ type: "tool_result", tool_use_id: block.id, content: danger, is_error: true });
|
|
continue;
|
|
}
|
|
try {
|
|
const output = await sshExec(input.command, input.timeout || 30);
|
|
toolResults.push({ type: "tool_result", tool_use_id: block.id, content: output || "(no output)" });
|
|
} catch (err: any) {
|
|
toolResults.push({ type: "tool_result", tool_use_id: block.id, content: `ERROR: ${err.message}`, is_error: true });
|
|
}
|
|
}
|
|
}
|
|
|
|
history.push({ role: "user", content: toolResults });
|
|
|
|
response = await anthropic.messages.create({
|
|
model: "claude-sonnet-4-20250514",
|
|
max_tokens: 1024,
|
|
system: SYSTEM_PROMPT,
|
|
tools,
|
|
messages: history,
|
|
});
|
|
}
|
|
|
|
// Extract text response
|
|
const text = response.content
|
|
.filter((b): b is Anthropic.TextBlock => b.type === "text")
|
|
.map((b) => b.text)
|
|
.join("\n");
|
|
|
|
history.push({ role: "assistant", content: text });
|
|
|
|
return text || "(no response)";
|
|
}
|
|
|
|
// ─── Security: block groups, unauthorized, rate-limited ───
|
|
bot.use(async (ctx, next) => {
|
|
// Block group/channel messages — private only
|
|
if (ctx.chat?.type !== "private") return;
|
|
|
|
// Silent block for unauthorized users — no response, no leak
|
|
if (!isAuthorized(ctx)) {
|
|
console.log(`🚫 Blocked user: ${ctx.from?.id} (@${ctx.from?.username})`);
|
|
return;
|
|
}
|
|
|
|
// Rate limit
|
|
if (ctx.from && isRateLimited(ctx.from.id)) {
|
|
return ctx.reply("⏳ Slow down — rate limited.");
|
|
}
|
|
|
|
await next();
|
|
});
|
|
|
|
// ─── Bot handlers ───
|
|
bot.command("start", (ctx) => {
|
|
return ctx.reply(
|
|
"🤖 *CR Management Bot*\n\n" +
|
|
"I'm connected to your server. You can:\n" +
|
|
"• Ask me anything about your infra\n" +
|
|
"• Ask me to run commands\n" +
|
|
"• `/status` — quick server health\n" +
|
|
"• `/clear` — reset conversation\n" +
|
|
"• `/ssh <command>` — run a command directly\n",
|
|
{ parse_mode: "Markdown" }
|
|
);
|
|
});
|
|
|
|
bot.command("status", async (ctx) => {
|
|
await ctx.reply("🔍 Checking server...");
|
|
try {
|
|
const output = await sshExec("uptime && echo '---' && df -h / && echo '---' && free -h | head -2 && echo '---' && incus list --format csv -c ns 2>/dev/null");
|
|
return ctx.reply(`\`\`\`\n${output}\n\`\`\``, { parse_mode: "Markdown" });
|
|
} catch (err: any) {
|
|
return ctx.reply(`❌ SSH failed: ${err.message}`);
|
|
}
|
|
});
|
|
|
|
bot.command("ssh", async (ctx) => {
|
|
const cmd = ctx.match;
|
|
if (!cmd) return ctx.reply("Usage: `/ssh <command>`", { parse_mode: "Markdown" });
|
|
const danger = isDangerous(cmd);
|
|
if (danger) return ctx.reply(danger);
|
|
await ctx.reply(`⚡ Running: \`${cmd}\``, { parse_mode: "Markdown" });
|
|
try {
|
|
const output = await sshExec(cmd);
|
|
return ctx.reply(`\`\`\`\n${output || "(no output)"}\n\`\`\``, { parse_mode: "Markdown" });
|
|
} catch (err: any) {
|
|
return ctx.reply(`❌ ${err.message}`);
|
|
}
|
|
});
|
|
|
|
bot.command("clear", (ctx) => {
|
|
conversations.delete(ctx.from!.id);
|
|
return ctx.reply("🧹 Conversation cleared.");
|
|
});
|
|
|
|
// Free text → Claude
|
|
bot.on("message:text", async (ctx) => {
|
|
const userId = ctx.from!.id;
|
|
const msg = ctx.message!.text;
|
|
|
|
await ctx.replyWithChatAction("typing");
|
|
|
|
try {
|
|
const reply = await chat(userId, msg);
|
|
console.log(`📤 [${userId}] reply length: ${reply.length}`);
|
|
// Split long messages (Telegram 4096 char limit)
|
|
const chunks = reply.match(/[\s\S]{1,4000}/g) || ["(empty)"];
|
|
for (const chunk of chunks) {
|
|
try {
|
|
await ctx.reply(chunk);
|
|
} catch (sendErr: any) {
|
|
console.error(`❌ [${userId}] send failed:`, sendErr.message);
|
|
await ctx.reply("⚠️ Failed to format response. Check logs.");
|
|
}
|
|
}
|
|
} catch (err: any) {
|
|
console.error(`❌ [${userId}]`, err);
|
|
await ctx.reply(`❌ Error: ${err.message}`);
|
|
}
|
|
});
|
|
|
|
// ─── Launch ───
|
|
console.log("🚀 Bot starting...");
|
|
bot.start({
|
|
onStart: (info) => console.log(`✅ @${info.username} is live`),
|
|
});
|