CookieBridge/projects/cookiebridge/src/relay/tokens.ts

import sodium from "sodium-native";
import type { DeviceInfo, AgentToken } from "../protocol/spec.js";

function generateToken(): string {
  const buf = Buffer.alloc(32);
  sodium.randombytes_buf(buf);
  return "cb_" + buf.toString("hex");
}

function generateId(): string {
  const buf = Buffer.alloc(16);
  sodium.randombytes_buf(buf);
  return buf.toString("hex");
}

/**
 * Manages device registrations and API tokens.
 */
export class DeviceRegistry {
  private devices = new Map<string, DeviceInfo>(); // deviceId -> info
  private tokenToDevice = new Map<string, string>(); // token -> deviceId
  // deviceId -> set of paired deviceIds
  private pairings = new Map<string, Set<string>>();

  register(deviceId: string, name: string, platform: string, encPub: string): DeviceInfo {
    const existing = this.devices.get(deviceId);
    if (existing) return existing;

    const token = generateToken();
    const info: DeviceInfo = {
      deviceId,
      name,
      platform,
      encPub,
      token,
      createdAt: new Date().toISOString(),
    };
    this.devices.set(deviceId, info);
    this.tokenToDevice.set(token, deviceId);
    return info;
  }

  getByToken(token: string): DeviceInfo | null {
    const deviceId = this.tokenToDevice.get(token);
    if (!deviceId) return null;
    return this.devices.get(deviceId) ?? null;
  }

  getById(deviceId: string): DeviceInfo | null {
    return this.devices.get(deviceId) ?? null;
  }

  /** Record that two devices are paired. */
  addPairing(deviceIdA: string, deviceIdB: string): void {
    let setA = this.pairings.get(deviceIdA);
    if (!setA) { setA = new Set(); this.pairings.set(deviceIdA, setA); }
    setA.add(deviceIdB);

    let setB = this.pairings.get(deviceIdB);
    if (!setB) { setB = new Set(); this.pairings.set(deviceIdB, setB); }
    setB.add(deviceIdA);
  }

  /** Get all deviceIds paired with a given device. */
  getPairedDevices(deviceId: string): string[] {
    const set = this.pairings.get(deviceId);
    return set ? Array.from(set) : [];
  }

  /** Get all deviceIds in the same pairing group (including self). */
  getPairingGroup(deviceId: string): string[] {
    const paired = this.getPairedDevices(deviceId);
    return [deviceId, ...paired];
  }
}

/**
 * Manages agent API tokens for the Agent Skill API.
 */
export class AgentRegistry {
  private agents = new Map<string, AgentToken>(); // id -> agent
  private tokenToAgent = new Map<string, string>(); // token -> agentId
  // agentId -> set of deviceIds that granted access
  private agentDeviceAccess = new Map<string, Set<string>>();

  create(name: string, encPub: string, allowedDomains: string[]): AgentToken {
    const id = generateId();
    const token = generateToken();
    const agent: AgentToken = {
      id,
      name,
      token,
      encPub,
      allowedDomains,
      createdAt: new Date().toISOString(),
    };
    this.agents.set(id, agent);
    this.tokenToAgent.set(token, id);
    return agent;
  }

  getByToken(token: string): AgentToken | null {
    const agentId = this.tokenToAgent.get(token);
    if (!agentId) return null;
    return this.agents.get(agentId) ?? null;
  }

  /** Grant an agent access to a device's cookies. */
  grantAccess(agentId: string, deviceId: string): void {
    let set = this.agentDeviceAccess.get(agentId);
    if (!set) { set = new Set(); this.agentDeviceAccess.set(agentId, set); }
    set.add(deviceId);
  }

  /** Get deviceIds an agent has access to. */
  getAccessibleDevices(agentId: string): string[] {
    const set = this.agentDeviceAccess.get(agentId);
    return set ? Array.from(set) : [];
  }

  revokeAccess(agentId: string, deviceId: string): void {
    const set = this.agentDeviceAccess.get(agentId);
    if (set) set.delete(deviceId);
  }
}