maestro/src/ssh/abuse-repo.ts

/**
 * SSH abuse counters — single source of truth for failure_count + lock_until.
 *
 * Design rationale (rev 4):
 *   Three scope kinds tracked per failure event, in one transaction:
 *     conn:<connId>           — always enforce
 *     userhost:<uid>|<host>|<user>  — always enforce
 *     globalhost:<host>|<user>      — enforce only for global connections;
 *                                     notification-only when paired with a
 *                                     user-owned connection (cross-user DoS
 *                                     mitigation without letting one user lock
 *                                     out another's globals).
 *
 *   `isLocked(connectionId)` reads ONLY the 'conn' scope row — abuse on a
 *   different user's connection to the same host should not block this one.
 *
 *   This module owns ALL failure_count / lock_until state. ssh_connections has
 *   no such columns; do not add them.
 *
 * Plan: docs/superpowers/plans/2026-05-12-ssh-tool-integration.md (Phase 1).
 */
import type Database from 'better-sqlite3';

export interface AbuseThresholds {
  windowMinutes: number;
  failureThreshold: number;
  lockMinutes: number;
}

export type AbuseScopeKind = 'conn' | 'userhost' | 'globalhost';

export interface RecordFailureArgs {
  connectionId: string;
  /** NULL for global connections; non-null for user-owned. */
  ownerId: string | null;
  userId: string;
  host: string;
  username: string;
  /** Defaults to new Date(). */
  now?: Date;
}

export interface RecordFailureResult {
  /** True if any enforced scope is now locked (after this update). */
  locked: boolean;
  /** First enforced scope that is locked, priority conn > userhost > globalhost. */
  lockedScope?: AbuseScopeKind;
  /** ISO8601 timestamp when the lock expires (from lockedScope row). */
  lockUntil?: string;
  /** True if globalhost scope just transitioned to threshold this call. */
  notifyAdmin?: boolean;
}

export interface IsLockedResult {
  locked: boolean;
  /** ISO8601 timestamp when the lock expires. */
  until?: string;
}

interface RawRow {
  scope_key: string;
  scope_kind: AbuseScopeKind;
  enforce_lock: number;
  failure_count: number;
  failure_window_start: string | null;
  lock_until: string | null;
  updated_at: string;
}

export interface SshAbuseRepo {
  checkAndRecordFailure(args: RecordFailureArgs): RecordFailureResult;
  isLocked(connectionId: string, now?: Date): IsLockedResult;
  /** Clear the 'conn' scope counter for a successful connection. */
  recordSuccess(connectionId: string, now?: Date): void;
  /** Admin: drop a single scope row (force unlock). Returns true if a row was removed. */
  reset(scopeKey: string): boolean;
  /** Inspect a single scope row (admin / tests). */
  getByScopeKey(scopeKey: string): RawRow | null;
}

function connKey(id: string) {
  return `conn:${id}`;
}
function userhostKey(userId: string, host: string, username: string) {
  return `userhost:${userId}|${host}|${username}`;
}
function globalhostKey(host: string, username: string) {
  return `globalhost:${host}|${username}`;
}

export function createAbuseRepo(
  db: Database.Database,
  thresholds: AbuseThresholds,
): SshAbuseRepo {
  if (thresholds.windowMinutes <= 0 || thresholds.failureThreshold <= 0 || thresholds.lockMinutes <= 0) {
    throw new Error('abuse: thresholds must be positive');
  }
  const windowMs = thresholds.windowMinutes * 60_000;
  const lockMs = thresholds.lockMinutes * 60_000;
  const threshold = thresholds.failureThreshold;

  const selectStmt = db.prepare(`SELECT * FROM ssh_abuse_counters WHERE scope_key = ?`);
  const insertStmt = db.prepare(`
    INSERT INTO ssh_abuse_counters
      (scope_key, scope_kind, enforce_lock, failure_count, failure_window_start, lock_until, updated_at)
    VALUES (?, ?, ?, ?, ?, ?, ?)
  `);
  const updateStmt = db.prepare(`
    UPDATE ssh_abuse_counters
    SET failure_count = ?,
        failure_window_start = ?,
        lock_until = ?,
        updated_at = ?
    WHERE scope_key = ?
  `);
  const deleteStmt = db.prepare(`DELETE FROM ssh_abuse_counters WHERE scope_key = ?`);

  interface ScopeOutcome {
    scopeKey: string;
    scopeKind: AbuseScopeKind;
    enforce: boolean;
    locked: boolean;
    lockUntil?: string;
    /** True if this call pushed count from <threshold to >=threshold. */
    justHitThreshold: boolean;
  }

  function applyToScope(
    scopeKey: string,
    scopeKind: AbuseScopeKind,
    enforce: boolean,
    nowIso: string,
    nowMs: number,
  ): ScopeOutcome {
    const row = selectStmt.get(scopeKey) as RawRow | undefined;
    const enforceFlag = enforce ? 1 : 0;
    let newCount = 1;
    let windowStart = nowIso;
    let lockUntil: string | null = null;
    let prevCount = 0;
    if (row) {
      prevCount = row.failure_count;
      const winStartMs = row.failure_window_start ? Date.parse(row.failure_window_start) : 0;
      const lockUntilMs = row.lock_until ? Date.parse(row.lock_until) : 0;
      // If currently locked, leave lock_until untouched; just bump count.
      if (lockUntilMs > nowMs) {
        newCount = row.failure_count + 1;
        windowStart = row.failure_window_start ?? nowIso;
        lockUntil = row.lock_until;
      } else if (winStartMs > 0 && nowMs - winStartMs < windowMs) {
        // Same window — increment.
        newCount = row.failure_count + 1;
        windowStart = row.failure_window_start ?? nowIso;
        if (newCount >= threshold) {
          lockUntil = new Date(nowMs + lockMs).toISOString();
        }
      } else {
        // Window expired or no prior window — fresh start.
        newCount = 1;
        windowStart = nowIso;
      }
      updateStmt.run(newCount, windowStart, lockUntil, nowIso, scopeKey);
    } else {
      insertStmt.run(scopeKey, scopeKind, enforceFlag, 1, nowIso, null, nowIso);
    }
    const justHit = prevCount < threshold && newCount >= threshold;
    return {
      scopeKey,
      scopeKind,
      enforce,
      locked: lockUntil != null && Date.parse(lockUntil) > nowMs,
      lockUntil: lockUntil ?? undefined,
      justHitThreshold: justHit,
    };
  }

  return {
    checkAndRecordFailure(args) {
      const now = args.now ?? new Date();
      const nowIso = now.toISOString();
      const nowMs = now.getTime();
      const isGlobal = args.ownerId === null;
      const tx = db.transaction((): RecordFailureResult => {
        const connOut = applyToScope(connKey(args.connectionId), 'conn', true, nowIso, nowMs);
        const userOut = applyToScope(
          userhostKey(args.userId, args.host, args.username),
          'userhost',
          true,
          nowIso,
          nowMs,
        );
        const globalOut = applyToScope(
          globalhostKey(args.host, args.username),
          'globalhost',
          isGlobal,
          nowIso,
          nowMs,
        );
        // Aggregate: first enforced + locked scope, priority conn > userhost > globalhost.
        const ordered = [connOut, userOut, globalOut];
        const lockedEnforced = ordered.find((o) => o.enforce && o.locked);
        const result: RecordFailureResult = { locked: false };
        if (lockedEnforced) {
          result.locked = true;
          result.lockedScope = lockedEnforced.scopeKind;
          result.lockUntil = lockedEnforced.lockUntil;
        }
        if (globalOut.justHitThreshold) {
          result.notifyAdmin = true;
        }
        return result;
      });
      return tx();
    },

    isLocked(connectionId, now) {
      const nowMs = (now ?? new Date()).getTime();
      const row = selectStmt.get(connKey(connectionId)) as RawRow | undefined;
      if (!row || row.lock_until == null) return { locked: false };
      if (Date.parse(row.lock_until) <= nowMs) return { locked: false };
      return { locked: true, until: row.lock_until };
    },

    recordSuccess(connectionId) {
      // Clearing the conn-scope counter on success lets a successful connection
      // 'forgive' prior failures within the window. userhost/globalhost remain
      // unchanged — they track host-level patterns, not per-connection state.
      deleteStmt.run(connKey(connectionId));
    },

    reset(scopeKey) {
      const r = deleteStmt.run(scopeKey);
      return r.changes > 0;
    },

    getByScopeKey(scopeKey) {
      const r = selectStmt.get(scopeKey) as RawRow | undefined;
      return r ?? null;
    },
  };
}