paif/auth.js

'use strict';

const fs = require('fs');
const path = require('path');
const os = require('os');
const { randomBytes } = require('crypto');

const BASE_DIR = path.join(os.homedir(), '.memory-bridge');
const ENV_PATH = path.join(BASE_DIR, '.env');
const AGENTS_DIR = path.join(BASE_DIR, 'agents');

let masterToken = null;

// Map of token -> agent_id (populated at init)
const tokenToAgent = new Map();

function init() {
  // Ensure agents directory exists
  fs.mkdirSync(AGENTS_DIR, { recursive: true });

  // Prefer env var (loaded by dotenv on subsequent runs)
  if (process.env.MEMORY_BRIDGE_TOKEN) {
    masterToken = process.env.MEMORY_BRIDGE_TOKEN;
  } else {
    // Try reading directly from file
    try {
      const raw = fs.readFileSync(ENV_PATH, 'utf8');
      const m = raw.match(/^MEMORY_BRIDGE_TOKEN=(.+)$/m);
      if (m) {
        masterToken = m[1].trim();
        process.env.MEMORY_BRIDGE_TOKEN = masterToken;
      }
    } catch { /* file doesn't exist yet */ }
  }

  // First run: generate and persist master token
  if (!masterToken) {
    masterToken = randomBytes(32).toString('hex');
    const dir = path.dirname(ENV_PATH);
    fs.mkdirSync(dir, { recursive: true });

    if (fs.existsSync(ENV_PATH)) {
      fs.appendFileSync(ENV_PATH, `\nMEMORY_BRIDGE_TOKEN=${masterToken}\n`);
    } else {
      fs.writeFileSync(
        ENV_PATH,
        `MEMORY_BRIDGE_TOKEN=${masterToken}\nMEMORY_BRIDGE_PORT=3722\nOLLAMA_URL=http://localhost:11434\n`
      );
    }

    process.env.MEMORY_BRIDGE_TOKEN = masterToken;

    console.log('\n══════════════════════════════════════════════════');
    console.log('First run — MASTER token (save this):');
    console.log(`  ${masterToken}`);
    console.log(`Persisted to: ${ENV_PATH}`);
    console.log('══════════════════════════════════════════════════\n');
  }

  // Load all agent tokens into memory
  loadAgentTokens();
}

function loadAgentTokens() {
  tokenToAgent.clear();
  try {
    const agents = fs.readdirSync(AGENTS_DIR).filter(d => {
      const stat = fs.statSync(path.join(AGENTS_DIR, d));
      return stat.isDirectory();
    });

    for (const agentId of agents) {
      const agentEnvPath = path.join(AGENTS_DIR, agentId, '.env');
      try {
        const raw = fs.readFileSync(agentEnvPath, 'utf8');
        const m = raw.match(/^AGENT_TOKEN=(.+)$/m);
        if (m) {
          const token = m[1].trim();
          tokenToAgent.set(token, agentId);
        }
      } catch {
        // Agent dir exists but no .env file
      }
    }
  } catch (err) {
    console.error('Failed to load agent tokens:', err.message);
  }
}

function getAgentIdFromToken(token) {
  // Master token has no agent_id (or special value)
  if (token === masterToken) {
    return null; // null = master/admin access
  }
  return tokenToAgent.get(token) || null;
}

function isMasterToken(token) {
  return token === masterToken;
}

function middleware(req, res, next) {
  const hdr = req.headers['authorization'] || '';
  if (!hdr.startsWith('Bearer ')) {
    return res.status(401).json({ error: 'Authorization: Bearer <token> required' });
  }

  const token = hdr.slice(7).trim();
  const agentId = getAgentIdFromToken(token);

  if (agentId === null && !isMasterToken(token)) {
    return res.status(401).json({ error: 'Invalid token' });
  }

  // Attach agent context to request
  req.agentContext = {
    token,
    agentId,        // null for master
    isMaster: isMasterToken(token)
  };

  next();
}

// Register a new agent (returns the agent token)
function registerAgent(agentId) {
  if (!/^[a-z0-9_-]+$/.test(agentId)) {
    throw new Error('Invalid agent_id: must be lowercase alphanumeric with _ or -');
  }

  const agentDir = path.join(AGENTS_DIR, agentId);
  const agentEnvPath = path.join(agentDir, '.env');

  if (fs.existsSync(agentEnvPath)) {
    // Agent already exists, return existing token
    const raw = fs.readFileSync(agentEnvPath, 'utf8');
    const m = raw.match(/^AGENT_TOKEN=(.+)$/m);
    if (m) {
      return m[1].trim();
    }
  }

  // Create new agent
  fs.mkdirSync(agentDir, { recursive: true });
  const agentToken = randomBytes(32).toString('hex');
  fs.writeFileSync(agentEnvPath, `AGENT_TOKEN=${agentToken}\nAGENT_ID=${agentId}\n`);

  // Reload tokens
  loadAgentTokens();

  return agentToken;
}

// List all registered agents
function listAgents() {
  try {
    return fs.readdirSync(AGENTS_DIR).filter(d => {
      const stat = fs.statSync(path.join(AGENTS_DIR, d));
      return stat.isDirectory() && fs.existsSync(path.join(AGENTS_DIR, d, '.env'));
    });
  } catch {
    return [];
  }
}

// Get agent's index directory path
function getAgentIndexDir(agentId) {
  if (!agentId) {
    throw new Error('agentId is required');
  }
  return path.join(BASE_DIR, 'indexes', agentId);
}

module.exports = {
  init,
  middleware,
  registerAgent,
  listAgents,
  getAgentIndexDir,
  isMaster: isMasterToken,
  getAgentIdFromToken
};