ShieldAI/packages/extension/src/content/index.ts

import { BackgroundMessage, MessageType, UrlCheckResult, UrlVerdict } from '../types';

let currentUrlVerdict: UrlVerdict | null = null;
let statusBar: HTMLElement | null = null;

chrome.runtime.onMessage.addListener(
  (message: BackgroundMessage) => {
    switch (message.type) {
      case MessageType.CHECK_URL_RESPONSE: {
        const result = message.payload as UrlCheckResult;
        currentUrlVerdict = result.verdict;
        updateStatusBar(result);
        injectPageBanner(result);
        highlightSuspiciousLinks(result);
        break;
      }
    }
  }
);

chrome.runtime.onInstalled.addListener(() => {
  chrome.storage.sync.get('shieldaiSettings', (data) => {
    if (data.shieldaiSettings?.enabled !== false) {
      requestUrlCheck();
    }
  });
});

function requestUrlCheck(): void {
  const url = window.location.href;
  if (url.startsWith('chrome://') || url.startsWith('chrome-extension://')) return;

  chrome.runtime.sendMessage({ type: MessageType.CHECK_URL, payload: { url } }).catch(() => {});
}

function updateStatusBar(result: UrlCheckResult): void {
  if (!statusBar) {
    statusBar = document.createElement('div');
    statusBar.id = 'shieldai-status-bar';
    Object.assign(statusBar.style, {
      position: 'fixed',
      top: '0',
      left: '0',
      right: '0',
      height: '3px',
      zIndex: '2147483647',
      transition: 'background-color 0.3s ease',
    });
    document.documentElement.insertBefore(statusBar, document.documentElement.firstChild);
  }

  const colors: Record<UrlVerdict, string> = {
    [UrlVerdict.SAFE]: '#22c55e',
    [UrlVerdict.SUSPICIOUS]: '#f59e0b',
    [UrlVerdict.PHISHING]: '#ef4444',
    [UrlVerdict.SPAM]: '#f97316',
    [UrlVerdict.EXPOSED_CREDENTIALS]: '#a855f7',
    [UrlVerdict.UNKNOWN]: '#6b7280',
  };

  statusBar.style.backgroundColor = colors[result.verdict] || colors[UrlVerdict.UNKNOWN];
  statusBar.title = `ShieldAI: ${result.verdict} (${result.threats.length} threat${result.threats.length !== 1 ? 's' : ''})`;
}

function injectPageBanner(result: UrlCheckResult): void {
  const existing = document.getElementById('shieldai-banner');
  if (existing) existing.remove();

  if (result.verdict === UrlVerdict.SAFE || result.verdict === UrlVerdict.UNKNOWN) return;

  const banner = document.createElement('div');
  banner.id = 'shieldai-banner';
  banner.innerHTML = `
    <div id="shieldai-banner-content">
      <span class="shieldai-icon">🛡️</span>
      <strong>ShieldAI:</strong> ${result.verdict.toUpperCase()} — ${result.threats[0]?.description || 'Potential threat detected'}
      <button id="shieldai-dismiss" style="margin-left: 12px; cursor: pointer; background: none; border: 1px solid #ccc; border-radius: 4px; padding: 2px 8px;">Dismiss</button>
    </div>
  `;

  Object.assign(banner.style, {
    position: 'fixed',
    top: '3px',
    left: '0',
    right: '0',
    zIndex: '2147483646',
    backgroundColor: result.verdict === UrlVerdict.PHISHING ? '#fef2f2' : '#fffbeb',
    borderBottom: `2px solid ${result.verdict === UrlVerdict.PHISHING ? '#ef4444' : '#f59e0b'}`,
    fontFamily: '-apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif',
    fontSize: '13px',
    color: '#374151',
  });

  const content = banner.querySelector('#shieldai-banner-content') as HTMLElement;
  Object.assign(content.style, {
    maxWidth: '800px',
    margin: '0 auto',
    padding: '8px 16px',
  });

  document.documentElement.insertBefore(banner, document.documentElement.firstChild.nextSibling);

  banner.querySelector('#shieldai-dismiss')?.addEventListener('click', () => {
    banner.remove();
  });
}

function highlightSuspiciousLinks(result: UrlCheckResult): void {
  if (result.verdict === UrlVerdict.SAFE) return;

  const links = document.querySelectorAll('a[href]');
  links.forEach((link) => {
    const href = link.getAttribute('href');
    if (!href) return;

    try {
      const linkDomain = new URL(href, window.location.href).hostname;
      const pageDomain = window.location.hostname;

      if (linkDomain !== pageDomain && !linkDomain.includes(pageDomain)) {
        link.classList.add('shieldai-external-link');
        link.title = `ShieldAI: External link → ${linkDomain}`;
      }
    } catch {
      // Relative or malformed URL
    }
  });

  const style = document.createElement('style');
  style.id = 'shieldai-link-styles';
  style.textContent = `
    a.shieldai-external-link::after {
      content: " ↗";
      opacity: 0.5;
      font-size: 0.8em;
    }
  `;
  document.head.appendChild(style);
}

requestUrlCheck();