Cookie Swapping & Session Profiles

import psl from 'psl';

function getRootDomain(hostname) {
  const parsed = psl.parse(hostname);
  return parsed.domain;
}

async function captureCookies(domain) {
  const rootDomain = getRootDomain(domain);
  const cookies = await chrome.cookies.getAll({ domain: rootDomain });
  const subdomainCookies = await chrome.cookies.getAll({ domain: '.' + rootDomain });

  // Merge and deduplicate
  const allCookies = [...cookies];
  for (const cookie of subdomainCookies) {
    const exists = allCookies.some(c =>
      c.name === cookie.name && c.domain === cookie.domain && c.path === cookie.path
    );
    if (!exists) {
      allCookies.push(cookie);
    }
  }
  return allCookies;
}

async function clearCookies(domain) {
  const cookies = await captureCookies(domain);
  for (const cookie of cookies) {
    const protocol = cookie.secure ? 'https://' : 'http://';
    const cookieDomain = cookie.domain.startsWith('.') 
      ? cookie.domain.slice(1) 
      : cookie.domain;
    const url = `${protocol}${cookieDomain}${cookie.path}`;
    await chrome.cookies.remove({ url, name: cookie.name });
  }
}

async function restoreCookies(cookies) {
  for (const cookie of cookies) {
    const { hostOnly, session, storeId, ...cookieData } = cookie;
    const protocol = cookie.secure ? 'https://' : 'http://';
    const cookieDomain = cookie.domain.startsWith('.') 
      ? cookie.domain.slice(1) 
      : cookie.domain;
    const url = `${protocol}${cookieDomain}${cookie.path}`;

    await chrome.cookies.set({
      url,
      ...cookieData,
      expirationDate: session ? undefined : cookieData.expirationDate
    });
  }
}

function showLoginRequiredNotification(domain, phase = 1) {
  const notification = document.createElement('div');
  notification.id = 'contextfort-login-notification';
  notification.style.cssText = `
    position: fixed;
    bottom: 20px;
    right: 20px;
    min-width: 360px;
    background: #1a1a1a;
    border-radius: 12px;
    padding: 20px;
    z-index: 2147483647;
  `;

  if (phase === 1) {
    // Phase 1: Initial prompt
    message.innerHTML = `
      No saved agent session for <strong>${domain}</strong><br><br>
      Your session has been cleared. You need to login manually for the agent to use this site.
    `;
    // Show "Ok, I'll log in" button
  } else {
    // Phase 2: After user acknowledges
    message.innerHTML = `
      Please login, then click "I've Logged In".<br><br>
      Or click "Restore My Session" to get your cookies back.
    `;
    // Show "I've Logged In" and "Restore My Session" buttons
  }
}

chrome.runtime.onMessage.addListener((message, sender, sendResponse) => {
  if (message.type === 'RESTORE_BUTTON_CLICKED') {
    const domain = message.domain;
    const tabId = sender.tab.id;

    (async () => {
      const humanSession = await loadSessionProfile(domain, 'human');
      if (humanSession) {
        await clearCookies(domain);
        await restoreCookies(humanSession.cookies);
        await chrome.tabs.reload(tabId);
      }
    })();
  }
});

async function swapAllDomainsToHuman(tabId, groupId) {
  console.log("swapping all domains to human");
  const swappedDomains = swappedDomainsPerSession.get(groupId);
  if (!swappedDomains || swappedDomains.size === 0) return true;
  
  for (const domain of swappedDomains) {
    // Save current agent session
    const agentSession = await captureFullSession(domain);
    if (agentSession) await saveSessionProfile(domain, 'agent', agentSession);
    
    // Restore human session
    const humanSession = await loadSessionProfile(domain, 'human');
    await clearCookies(domain);
    if (humanSession) await restoreCookies(humanSession.cookies);
  }

  // Reload current tab if it's on a swapped domain
  const tab = await chrome.tabs.get(tabId);
  const currentDomain = getRootDomain(new URL(tab.url).hostname);
  if (currentDomain && swappedDomains.has(currentDomain)) {
    await chrome.tabs.reload(tabId);
  }
  
  swappedDomainsPerSession.delete(groupId);
  return true;
}

const swappedDomainsPerSession = new Map();
// Example:
// swappedDomainsPerSession = Map {
//   123 => Set { "example.com", "api.example.com", "docs.example.com" },
//   456 => Set { "github.com" }
// }