Fix : MITM

src/mitm/server.js

@@ -3,6 +3,7 @@ const fs = require("fs");
 const path = require("path");
 const dns = require("dns");
 const { promisify } = require("util");
+const { log, err } = require("./logger");
 
 // Allow self-signed certs from MITM root CA when fetching external hosts
 
@@ -25,7 +26,7 @@ const DB_FILE = path.join(DATA_DIR, "db.json");
 const ENABLE_FILE_LOG = false;
 
 if (!API_KEY) {
-  console.error("❌ ROUTER_API_KEY required");
+  err("ROUTER_API_KEY required");
   process.exit(1);
 }
 
@@ -57,11 +58,11 @@ function sniCallback(servername, cb) {
 
     // Cache it
     certCache.set(servername, ctx);
-    console.log(`✅ Generated cert for: ${servername}`);
+    log(`🔐 Cert generated: ${servername}`);
 
     cb(null, ctx);
   } catch (error) {
-    console.error(`❌ SNI error for ${servername}:`, error.message);
+    err(`SNI error for ${servername}: ${error.message}`);
     cb(error);
   }
 }
@@ -79,7 +80,7 @@ try {
     SNICallback: sniCallback
   };
 } catch (e) {
-  console.error(`❌ Root CA not found in ${certDir}: ${e.message}`);
+  err(`Root CA not found in ${certDir}: ${e.message}`);
   process.exit(1);
 }
 
@@ -134,7 +135,13 @@ function getMappedModel(tool, model) {
   try {
     if (!fs.existsSync(DB_FILE)) return null;
     const db = JSON.parse(fs.readFileSync(DB_FILE, "utf-8"));
-    return db.mitmAlias?.[tool]?.[model] || null;
+    const aliases = db.mitmAlias?.[tool];
+    if (!aliases) return null;
+    // Exact match first
+    if (aliases[model]) return aliases[model];
+    // Prefix match fallback: find alias key that starts with model or model starts with key
+    const prefixKey = Object.keys(aliases).find(k => k && aliases[k] && (model.startsWith(k) || k.startsWith(model)));
+    return prefixKey ? aliases[prefixKey] : null;
   } catch {
     return null;
   }
@@ -167,8 +174,8 @@ async function passthrough(req, res, bodyBuffer) {
     forwardRes.pipe(res);
   });
 
-  forwardReq.on("error", (err) => {
-    console.error(`❌ Passthrough error: ${err.message}`);
+  forwardReq.on("error", (e) => {
+    err(`Passthrough error: ${e.message}`);
     if (!res.headersSent) res.writeHead(502);
     res.end("Bad Gateway");
   });
@@ -216,7 +223,7 @@ async function intercept(req, res, bodyBuffer, mappedModel) {
       res.write(decoder.decode(value, { stream: true }));
     }
   } catch (error) {
-    console.error(`❌ ${error.message}`);
+    err(`Intercept error: ${error.message}`);
     if (!res.headersSent) res.writeHead(500, { "Content-Type": "application/json" });
     res.end(JSON.stringify({ error: { message: error.message, type: "mitm_error" } }));
   }
@@ -250,30 +257,34 @@ const server = https.createServer(sslOptions, async (req, res) => {
     if (!isChat) return passthrough(req, res, bodyBuffer);
 
     const model = extractModel(req.url, bodyBuffer);
-    console.log("Extracted model:", model);
+    log(`🔍 model="${model}" url=${req.url}`);
     const mappedModel = getMappedModel(tool, model);
 
-    if (!mappedModel) return passthrough(req, res, bodyBuffer);
+    if (!mappedModel) {
+      log(`⏩ passthrough | no mapping | ${tool} | ${model || "unknown"}`);
+      return passthrough(req, res, bodyBuffer);
+    }
 
+    log(`⚡ intercept | ${tool} | ${model} → ${mappedModel}`);
     return intercept(req, res, bodyBuffer, mappedModel);
-  } catch (err) {
-    console.error(`❌ Unhandled request error: ${err.message}`);
+  } catch (e) {
+    err(`Unhandled request error: ${e.message}`);
     if (!res.headersSent) res.writeHead(500, { "Content-Type": "application/json" });
-    res.end(JSON.stringify({ error: { message: err.message, type: "mitm_error" } }));
+    res.end(JSON.stringify({ error: { message: e.message, type: "mitm_error" } }));
   }
 });
 
 server.listen(LOCAL_PORT, () => {
-  console.log(`🚀 MITM ready on :${LOCAL_PORT}`);
+  log(`🚀 Server ready on :${LOCAL_PORT}`);
 });
 
 server.on("error", (error) => {
   if (error.code === "EADDRINUSE") {
-    console.error(`❌ Port ${LOCAL_PORT} already in use`);
+    err(`Port ${LOCAL_PORT} already in use`);
   } else if (error.code === "EACCES") {
-    console.error(`❌ Permission denied for port ${LOCAL_PORT}`);
+    err(`Permission denied for port ${LOCAL_PORT}`);
   } else {
-    console.error(`❌ ${error.message}`);
+    err(error.message);
   }
   process.exit(1);
 });