dApp Browser Security

import { ThreatClient, Chain, type ScanChainId } from '@webacy-xyz/sdk';

const client = new ThreatClient({
  apiKey: process.env.WEBACY_API_KEY,
  defaultChain: Chain.ETH,
});

// Types
type RiskLevel = 'safe' | 'warning' | 'danger' | 'blocked';

interface UrlCheckResult {
  allowed: boolean;
  riskLevel: RiskLevel;
  message?: string;
}

interface TransactionCheckResult {
  allowed: boolean;
  riskLevel: RiskLevel;
  preview: {
    assetChanges: Array<{
      type: string;
      symbol: string;
      amount: string;
    }>;
    warnings: string[];
  };
  requiresConfirmation: boolean;
}

interface ConnectionInfo {
  domain: string;
  riskLevel: RiskLevel;
  connectedAt: Date;
  lastTransaction?: Date;
}

// URL Security
async function checkUrlSafety(url: string): Promise<UrlCheckResult> {
  const result = await client.url.check(url);

  switch (result.riskLevel) {
    case 'high':
      // Known phishing or blocklisted site
      return {
        allowed: false,
        riskLevel: 'blocked',
        message: `${result.description} Do not connect your wallet.`,
      };
    case 'medium':
      // Similar to a known project's domain: possible impersonation
      return {
        allowed: true,
        riskLevel: 'danger',
        message: `${result.description} Check the address bar before connecting.`,
      };
    case 'low':
      return { allowed: true, riskLevel: 'safe' };
    case 'unknown':
    default:
      // Fail toward caution: an unexpected level is never treated as safe.
      return {
        allowed: true,
        riskLevel: 'warning',
        message: 'This site has not been verified. Exercise caution.',
      };
  }
}

// Transaction Security
async function checkTransaction(
  userAddress: string,
  // `raw` = serialized unsigned tx (ethers `Transaction.unsignedSerialized`); `to` for the contract pre-check
  txData: { to: string; raw: string },
  chainId: ScanChainId,
  dappDomain: string
): Promise<TransactionCheckResult> {
  // First, verify the contract
  const chain = chainIdToChain(chainId);
  const contractAnalysis = await client.addresses.analyze(txData.to, { chain });

  // Check for known drainers
  const isDrainer = contractAnalysis.issues?.some(i =>
    i.tag.toLowerCase().includes('drainer')
  );

  if (isDrainer) {
    return {
      allowed: false,
      riskLevel: 'blocked',
      preview: {
        assetChanges: [],
        warnings: ['This contract has been identified as a wallet drainer.'],
      },
      requiresConfirmation: false,
    };
  }

  // Simulate the transaction
  const result = (await client.scan.scanTransaction(userAddress, {
    tx: { from: userAddress, raw: txData.raw },
    chain: chainId,
    domain: dappDomain,
  })) as unknown as {
    simulation: Array<{
      txData: { changeType: 'TRANSFER' | 'APPROVE'; rawAmount?: string; symbol?: string; assetType?: string; counterpartyAddress: string };
      counterpartyRisk: { high?: number; medium?: number; overallRisk?: number; issues?: Array<{ tags: Array<{ name: string }> }> };
      assetRisk: { high?: number; medium?: number };
      functionRisk?: { functionName: string; description: string };
    }>;
  };

  // Determine risk level from the recipient / spender and token profiles.
  // An empty simulation, or a recipient that could not be scored, is inconclusive.
  let riskLevel: RiskLevel = 'safe';
  const parties = result.simulation.flatMap((i) => [i.counterpartyRisk, i.assetRisk]);
  if (parties.some((p) => (p.high ?? 0) > 0)) {
    riskLevel = 'blocked';
  } else if (parties.some((p) => (p.medium ?? 0) > 0)) {
    riskLevel = 'danger';
  } else if (
    result.simulation.length === 0 ||
    result.simulation.some((i) => i.counterpartyRisk.overallRisk === undefined)
  ) {
    riskLevel = 'warning';
  }

  return {
    allowed: riskLevel !== 'blocked',
    riskLevel,
    preview: {
      assetChanges: result.simulation.map((i) => ({
        type: i.txData.changeType,
        symbol: i.txData.symbol ?? i.txData.assetType ?? 'asset',
        amount: i.txData.rawAmount ?? '',
      })),
      warnings: result.simulation.flatMap((i) => [
        ...(i.counterpartyRisk.issues?.flatMap((issue) => issue.tags.map((t) => t.name)) ?? []),
        ...(i.functionRisk ? [`Risky function: ${i.functionRisk.functionName}`] : []),
      ]),
    },
    requiresConfirmation: riskLevel === 'danger' || riskLevel === 'warning',
  };
}

// Connection Management
class ConnectionManager {
  private connections: Map<string, ConnectionInfo> = new Map();

  async connect(domain: string): Promise<{
    success: boolean;
    riskLevel: RiskLevel;
    message?: string;
  }> {
    const urlCheck = await checkUrlSafety(`https://${domain}`);

    if (!urlCheck.allowed) {
      return {
        success: false,
        riskLevel: urlCheck.riskLevel,
        message: urlCheck.message,
      };
    }

    this.connections.set(domain, {
      domain,
      riskLevel: urlCheck.riskLevel,
      connectedAt: new Date(),
    });

    return {
      success: true,
      riskLevel: urlCheck.riskLevel,
      message: urlCheck.message,
    };
  }

  disconnect(domain: string): void {
    this.connections.delete(domain);
  }

  getConnections(): ConnectionInfo[] {
    return Array.from(this.connections.values());
  }

  async refreshRiskLevels(): Promise<ConnectionInfo[]> {
    const flagged: ConnectionInfo[] = [];

    for (const [domain, info] of this.connections) {
      const { riskLevel } = await checkUrlSafety(`https://${domain}`);
      const previous = info.riskLevel;
      info.riskLevel = riskLevel;
      // Flag blocked/danger, and a site that was safe and is now unverified.
      if (
        riskLevel === 'blocked' ||
        riskLevel === 'danger' ||
        (riskLevel === 'warning' && previous === 'safe')
      ) {
        flagged.push(info);
      }
    }

    return flagged;
  }
}

// Browser Security Controller
class DappBrowserSecurity {
  private connectionManager = new ConnectionManager();

  // Call when user navigates to a new URL
  async onNavigate(url: string): Promise<UrlCheckResult> {
    return checkUrlSafety(url);
  }

  // Call when dApp requests wallet connection
  async onConnectionRequest(domain: string) {
    return this.connectionManager.connect(domain);
  }

  // Call when dApp requests transaction signature
  async onTransactionRequest(
    userAddress: string,
    txData: { to: string; raw: string },
    chainId: ScanChainId,
    dappDomain: string
  ) {
    return checkTransaction(userAddress, txData, chainId, dappDomain);
  }

  // Call periodically to check connected dApps
  async refreshConnections() {
    return this.connectionManager.refreshRiskLevels();
  }

  // Get all active connections
  getActiveConnections() {
    return this.connectionManager.getConnections();
  }

  // Disconnect a dApp
  disconnect(domain: string) {
    this.connectionManager.disconnect(domain);
  }
}

// Helper function
function chainIdToChain(chainId: ScanChainId): Chain {
  const mapping: Record<number, Chain> = {
    1: Chain.ETH,
    56: Chain.BSC,
    137: Chain.POL,
    10: Chain.OPT,
    42161: Chain.ARB,
    8453: Chain.BASE,
  };
  return mapping[chainId] ?? Chain.ETH;
}

// Export the controller
export const browserSecurity = new DappBrowserSecurity();