AeroFetch/src/main/poToken.ts

import { BrowserWindow, type WebContents } from 'electron'
import { LOGIN_PARTITION } from './cookies'

/** A public YouTube video used as the extraction target (stable, always accessible). */
const YT_VIDEO = 'https://www.youtube.com/watch?v=jNQXAC9IVRw'

/**
 * Restrict a PO-token window to *.youtube.com and accounts.google.com (for
 * sign-in). Mirrors hardenLoginWebContents in cookies.ts.
 */
function isAllowedPoTokenUrl(target: string): boolean {
  try {
    const { protocol, hostname } = new URL(target)
    if (protocol !== 'http:' && protocol !== 'https:') return false
    return (
      hostname === 'youtube.com' ||
      hostname.endsWith('.youtube.com') ||
      hostname === 'accounts.google.com' ||
      hostname.endsWith('.accounts.google.com')
    )
  } catch {
    return false
  }
}

function hardenPoTokenWebContents(wc: WebContents): void {
  wc.setWindowOpenHandler((details) => {
    if (!isAllowedPoTokenUrl(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 (!isAllowedPoTokenUrl(navUrl)) e.preventDefault()
  })
  wc.on('will-redirect', (e, navUrl) => {
    if (!isAllowedPoTokenUrl(navUrl)) e.preventDefault()
  })
  wc.on('did-create-window', (child) => hardenPoTokenWebContents(child.webContents))
}

/**
 * Async IIFE injected into the YouTube page after load. Polls
 * ytInitialPlayerResponse.serviceIntegrityDimensions.poToken (the field
 * yt-dlp's youtube extractor reads) for up to 8 s, returning the token
 * string or null if it never appears. The poll loop is needed because the
 * player JS initialises asynchronously after the DOM is ready.
 */
const EXTRACT_SCRIPT = `
(async function() {
  for (let i = 0; i < 16; i++) {
    try {
      const pot = window.ytInitialPlayerResponse
        && window.ytInitialPlayerResponse.serviceIntegrityDimensions
        && window.ytInitialPlayerResponse.serviceIntegrityDimensions.poToken
      if (typeof pot === 'string' && pot.length > 0) return pot
    } catch (_) {}
    await new Promise(r => setTimeout(r, 500))
  }
  return null
})()
`

let mintWindow: BrowserWindow | null = null

/**
 * Open a visible BrowserWindow on a YouTube video page and extract the
 * Proof-of-Origin token from the page runtime. The window uses the shared
 * login session so a signed-in user's credentials are available.
 *
 * Resolves with the token string on success, or null if the window was closed
 * by the user before extraction completed or if the field was not found.
 */
export function openPoTokenWindow(
  /** app-theme window background (UI23) — supplied by the IPC layer so this
   *  window's chrome matches the app's Light/Dark instead of flashing white. */
  backgroundColor?: string
): Promise<string | null> {
  // If a minting window is already open, bring it to the front rather than
  // opening a second one.
  if (mintWindow && !mintWindow.isDestroyed()) {
    mintWindow.focus()
    // Return a promise that resolves when the existing window closes.
    return new Promise((resolve) => {
      mintWindow!.once('closed', () => resolve(null))
    })
  }

  return new Promise((resolve) => {
    let resolved = false
    const finish = (token: string | null): void => {
      if (resolved) return
      resolved = true
      resolve(token)
    }

    let win: BrowserWindow
    try {
      win = new BrowserWindow({
        width: 960,
        height: 640,
        title: 'AeroFetch — Fetch YouTube token',
        // Match the app's resolved theme (UI23), like the cookie sign-in window.
        backgroundColor,
        autoHideMenuBar: true,
        webPreferences: {
          partition: LOGIN_PARTITION,
          sandbox: true,
          contextIsolation: true,
          nodeIntegration: false
        }
      })
    } catch {
      finish(null)
      return
    }

    mintWindow = win
    hardenPoTokenWebContents(win.webContents)
    win.webContents.session.setPermissionRequestHandler((_wc, _perm, cb) => cb(false))

    win.webContents.once('did-finish-load', () => {
      win.webContents
        .executeJavaScript(EXTRACT_SCRIPT, true)
        .then((result: unknown) => {
          const token = typeof result === 'string' && result.length > 0 ? result : null
          finish(token)
          win.close()
        })
        .catch(() => {
          finish(null)
          win.close()
        })
    })

    win.on('closed', () => {
      mintWindow = null
      finish(null)
    })

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