AeroFetch/src/main/cookies.ts

import { app, safeStorage, session, BrowserWindow, type Cookie, type WebContents } from 'electron'
import { existsSync, statSync, unlinkSync, writeFileSync, readFileSync } from 'fs'
import { join } from 'path'
import { assertHttpUrl } from './url'
import type { CookiesStatus, CookiesLoginResult } from '@shared/ipc'

/**
 * Persisted session AeroFetch's built-in sign-in window uses. Kept separate
 * from the main window's (default) session so a logged-in site can't see or
 * touch anything the app itself loads. Exported so the PO-token window can
 * share the same session (and thus the user's YouTube sign-in state).
 */
export const LOGIN_PARTITION = 'persist:aerofetch-login'

/**
 * The sign-in window renders untrusted remote content, so every navigation and
 * popup is confined to web URLs — http(s), plus about:blank. This stops a
 * logged-in (or malicious) page from steering the window to a file:// URL, to
 * the app's own aerofetch:// protocol handler, or to any other external URI
 * scheme used as a pivot. (audit T4)
 */
export function isAllowedLoginUrl(target: string): boolean {
  try {
    const { protocol } = new URL(target)
    return protocol === 'http:' || protocol === 'https:' || protocol === 'about:'
  } catch {
    return false
  }
}

/**
 * Apply the web-only confinement to a sign-in webContents: block top-frame
 * navigations and server redirects to non-web URLs, restrict popups to web URLs
 * (reusing the same secure, partitioned webPreferences), and recurse into any
 * popup the page opens so a nested window can't escape the policy either. The
 * window-open handler alone only governs NEW windows, not navigations of an
 * existing one — both vectors are covered here. (audit T4)
 */
function hardenLoginWebContents(wc: WebContents): void {
  wc.setWindowOpenHandler((details) => {
    if (!isAllowedLoginUrl(details.url)) return { action: 'deny' }
    return {
      action: 'allow',
      overrideBrowserWindowOptions: {
        autoHideMenuBar: true,
        webPreferences: {
          partition: LOGIN_PARTITION,
          sandbox: true,
          contextIsolation: true,
          nodeIntegration: false
        }
      }
    }
  })
  wc.on('will-navigate', (e, navUrl) => {
    if (!isAllowedLoginUrl(navUrl)) e.preventDefault()
  })
  wc.on('will-redirect', (e, navUrl) => {
    if (!isAllowedLoginUrl(navUrl)) e.preventDefault()
  })
  wc.on('did-create-window', (child) => hardenLoginWebContents(child.webContents))
}

// The persisted cookie jar is encrypted at rest (H7). yt-dlp's `--cookies` needs
// a plaintext file, so the stored form is encrypted via Electron safeStorage
// (DPAPI on Windows — the same protection settings secrets get) and only ever
// decrypted to a short-lived temp file for the duration of a download. This
// matters for the portable build, where userData sits next to the exe (USB /
// shared Downloads folder) and a plaintext jar would hand over a logged-in session.
function cookieStorePath(): string {
  return join(app.getPath('userData'), 'cookies.dat')
}

function legacyCookiePath(): string {
  return join(app.getPath('userData'), 'cookies.txt')
}

// Marks ciphertext so a read can tell it from legacy/fallback plaintext (written
// before encryption, or where safeStorage was unavailable). Mirrors settings.ts.
const ENC_PREFIX = 'enc:v1:'

/** Persist the Netscape cookie text, encrypted where safeStorage is available. */
function storeCookies(netscape: string): void {
  let out = netscape
  try {
    if (safeStorage.isEncryptionAvailable()) {
      out = ENC_PREFIX + safeStorage.encryptString(netscape).toString('base64')
    }
  } catch {
    /* fall through — store plaintext, as it was before encryption existed */
  }
  writeFileSync(cookieStorePath(), out)
}

/** Decrypt the persisted cookie text, or null if none / undecryptable here. */
function loadCookies(): string | null {
  const p = cookieStorePath()
  if (!existsSync(p)) return null
  const stored = readFileSync(p, 'utf8')
  if (!stored.startsWith(ENC_PREFIX)) return stored // legacy / fallback plaintext
  try {
    return safeStorage.decryptString(Buffer.from(stored.slice(ENC_PREFIX.length), 'base64'))
  } catch {
    // A different Windows user/machine (portable jar copied elsewhere) or a
    // corrupt blob — unusable; treat as no cookies rather than feed ciphertext.
    return null
  }
}

/**
 * One-time migration of a pre-encryption plaintext `cookies.txt` into the
 * encrypted store, deleting the plaintext so a logged-in session no longer sits
 * in the open after an update (H7). Best-effort; safe to call on every launch.
 */
export function migrateLegacyCookies(): void {
  const legacy = legacyCookiePath()
  try {
    if (!existsSync(legacy)) return
    if (!existsSync(cookieStorePath())) storeCookies(readFileSync(legacy, 'utf8'))
    unlinkSync(legacy)
  } catch {
    /* best-effort */
  }
}

/** Whether a stored cookie jar exists (regardless of decryptability here). */
export function hasStoredCookies(): boolean {
  return existsSync(cookieStorePath())
}

/**
 * Decrypt the jar to a plaintext file at `dest` for yt-dlp's `--cookies`. Returns
 * false when there's nothing usable to write. The caller deletes `dest` once the
 * download settles, keeping the plaintext window as short as possible.
 */
export function materializeCookies(dest: string): boolean {
  const text = loadCookies()
  if (text == null) return false
  try {
    writeFileSync(dest, text)
    return true
  } catch {
    return false
  }
}

export function getCookiesStatus(): CookiesStatus {
  const p = cookieStorePath()
  if (!existsSync(p)) return { exists: false }
  return { exists: true, savedAt: statSync(p).mtimeMs }
}

export async function clearCookies(): Promise<void> {
  const p = cookieStorePath()
  if (existsSync(p)) unlinkSync(p)
  await session.fromPartition(LOGIN_PARTITION).clearStorageData()
}

/**
 * Render cookies in the Netscape cookie-jar format yt-dlp's `--cookies` reads.
 * Columns: domain, includeSubdomains, path, secure, expiry (unix seconds, 0
 * for session cookies), name, value.
 */
function toNetscapeCookieFile(cookies: Cookie[]): string {
  const lines = ['# Netscape HTTP Cookie File', '# Generated by AeroFetch — do not edit.', '']
  for (const c of cookies) {
    if (!c.domain || !c.name) continue
    const domain = c.hostOnly || c.domain.startsWith('.') ? c.domain : `.${c.domain}`
    const includeSubdomains = domain.startsWith('.') ? 'TRUE' : 'FALSE'
    const expiry = c.session || !c.expirationDate ? 0 : Math.round(c.expirationDate)
    lines.push(
      [
        domain,
        includeSubdomains,
        c.path || '/',
        c.secure ? 'TRUE' : 'FALSE',
        expiry,
        c.name,
        c.value
      ].join('\t')
    )
  }
  return lines.join('\n') + '\n'
}

let loginWindow: BrowserWindow | null = null
let pendingResolvers: Array<(r: CookiesLoginResult) => void> = []

function exportAndResolve(): void {
  session
    .fromPartition(LOGIN_PARTITION)
    .cookies.get({})
    .then((cookies) => {
      // Don't persist an empty jar: closing the window without signing in would
      // otherwise leave a useless "saved" cookie file behind (L50). The renderer
      // turns cookieCount === 0 into a "did you sign in?" hint.
      if (cookies.length > 0) storeCookies(toNetscapeCookieFile(cookies))
      const result: CookiesLoginResult = { ok: true, cookieCount: cookies.length }
      pendingResolvers.forEach((resolve) => resolve(result))
    })
    .catch((e) => {
      const result: CookiesLoginResult = { ok: false, error: (e as Error).message }
      pendingResolvers.forEach((resolve) => resolve(result))
    })
    .finally(() => {
      pendingResolvers = []
    })
}

/**
 * Open (or focus + navigate) a sign-in window backed by a persisted session
 * partition. The user logs into whatever site they need, then closes the
 * window — cookies are exported to a Netscape-format file at that point,
 * ready for yt-dlp's `--cookies`. Mirrors Seal's "log in via WebView" feature.
 */
export function openCookieLoginWindow(
  url: string,
  parent?: BrowserWindow
): Promise<CookiesLoginResult> {
  return new Promise((resolve) => {
    let validUrl: string
    try {
      validUrl = assertHttpUrl(url)
    } catch (e) {
      resolve({ ok: false, error: (e as Error).message })
      return
    }

    if (loginWindow && !loginWindow.isDestroyed()) {
      pendingResolvers.push(resolve)
      loginWindow
        .loadURL(validUrl)
        .catch((e) => console.error('[AeroFetch] cookie login loadURL failed:', e))
      loginWindow.focus()
      return
    }

    pendingResolvers = [resolve]
    let win: BrowserWindow
    try {
      win = new BrowserWindow({
        width: 480,
        height: 720,
        title: 'Sign in — AeroFetch',
        autoHideMenuBar: true,
        // Group under the app window instead of taking its own taskbar button
        // (W6); a child window stays above its parent without a modal block.
        parent: parent && !parent.isDestroyed() ? parent : undefined,
        webPreferences: {
          partition: LOGIN_PARTITION,
          sandbox: true,
          contextIsolation: true,
          nodeIntegration: false
        }
      })
    } catch (e) {
      pendingResolvers = []
      resolve({ ok: false, error: (e as Error).message })
      return
    }
    loginWindow = win

    // Confine every navigation/popup to web URLs (recursively, so OAuth/SSO
    // popups sharing the cookie partition are covered too), and deny all gated
    // web permissions — signing in needs no camera/mic/geolocation/etc., and the
    // page is untrusted. (audit T4)
    hardenLoginWebContents(win.webContents)
    win.webContents.session.setPermissionRequestHandler((_wc, _permission, cb) => cb(false))

    // 'close' fires on a normal user close; a programmatic destroy() fires only
    // 'closed'. Latch whichever runs first so the partition is exported exactly
    // once and the promise can never hang (B7) — without this, a destroy()
    // without a preceding 'close' would leave pendingResolvers pending forever.
    let exportStarted = false
    const exportOnce = (): void => {
      if (exportStarted) return
      exportStarted = true
      exportAndResolve()
    }
    win.on('closed', () => {
      loginWindow = null
      // Fallback for destroy()/closed-without-close: the partition outlives the
      // window, so a late export still captures whatever cookies were collected.
      exportOnce()
    })
    // Closing the window IS "I'm done" — export whatever the partition
    // collected. The partition itself outlives the window, so this is safe
    // even though cookies.get() resolves after 'close' has already fired.
    win.on('close', exportOnce)

    win.loadURL(validUrl).catch(() => {
      /* navigation errors surface as Chromium's own error page */
    })
  })
}