Bifrost – Automated virtual key generations script

<#
.SYNOPSIS
  cowork-login: ensure a Bifrost virtual key exists for this user.
  If missing, opens the Bifrost app for Okta sign-in, mints a VK via the
  governance API using the user's own session JWT, stores it locally, and
  emits the env vars Claude Cowork needs.

.USAGE
  # dot-source so the env vars land in your current session:
  . .\cowork-login.ps1
  # then launch Claude Cowork from the same shell.
#>

$ErrorActionPreference = "Stop"

# ---------------- Config ----------------
$Bifrost  = "https://bifrost.yourdomain.com"
$ConfDir  = Join-Path $env:USERPROFILE ".config\cowork-bifrost"
$VkFile   = Join-Path $ConfDir "vk"
$Port     = 8899
$AllowedModels = @("anthropic.claude-sonnet-4-6-v1:0")
$MonthlyBudget = 50.00
# ----------------------------------------

function Test-HaveVk {
    (Test-Path $VkFile) -and ((Get-Content $VkFile -Raw).Trim() -like "sk-bf-*")
}

function Invoke-BifrostApi {
    param(
        [string]$Path,
        [string]$Jwt,
        [string]$Method = "GET",
        [object]$Body = $null
    )
    $params = @{
        Uri     = "$Bifrost$Path"
        Method  = $Method
        Headers = @{ Authorization = "Bearer $Jwt" }
        ContentType = "application/json"
    }
    if ($Body) { $params.Body = ($Body | ConvertTo-Json -Depth 10) }
    Invoke-RestMethod @params
}

function Get-JwtEmail {
    param([string]$Jwt)
    $payload = $Jwt.Split(".")[1].Replace('-', '+').Replace('_', '/')
    switch ($payload.Length % 4) { 2 { $payload += "==" } 3 { $payload += "=" } }
    $json = [Text.Encoding]::UTF8.GetString([Convert]::FromBase64String($payload))
    ($json | ConvertFrom-Json).email
}

function Get-SessionJwt {
    <# Opens the browser to Bifrost login; captures the session JWT via a
       localhost HTTP listener. The callback page fetches the Bifrost session
       token (cookie-authenticated) and POSTs it back to us, then closes. #>
    $state    = [Guid]::NewGuid().ToString("N")
    $listener = [System.Net.HttpListener]::new()
    $listener.Prefixes.Add("http://127.0.0.1:$Port/")
    $listener.Start()

    Start-Process "$Bifrost/login?redirect=http://localhost:$Port/callback"
    Write-Host "Waiting for Okta sign-in in your browser..." -ForegroundColor Cyan

    $jwt      = $null
    $deadline = (Get-Date).AddMinutes(5)

    $callbackHtml = @"
<!doctype html><html><body><h3>Finishing sign-in...</h3>
<script>
  fetch('$Bifrost/api/auth/session', {credentials:'include'})
    .then(r => r.json())
    .then(d => fetch('/token?state=$state', {method:'POST', body:d.token}))
    .then(() => { document.body.innerHTML='<h3>Done — you can close this tab.</h3>'; window.close(); })
    .catch(e => { document.body.innerHTML='<h3>Error: '+e+'</h3>'; });
</script></body></html>
"@

    try {
        while (-not $jwt -and (Get-Date) -lt $deadline) {
            $ctxTask = $listener.GetContextAsync()
            while (-not $ctxTask.AsyncWaitHandle.WaitOne(500)) {
                if ((Get-Date) -ge $deadline) { throw "Timed out waiting for login." }
            }
            $ctx = $ctxTask.Result
            $req = $ctx.Request
            $res = $ctx.Response

            if ($req.HttpMethod -eq "GET" -and $req.Url.AbsolutePath -eq "/callback") {
                $bytes = [Text.Encoding]::UTF8.GetBytes($callbackHtml)
                $res.ContentType = "text/html"
                $res.OutputStream.Write($bytes, 0, $bytes.Length)
                $res.Close()
            }
            elseif ($req.HttpMethod -eq "POST" -and
                    $req.Url.AbsolutePath -eq "/token" -and
                    $req.QueryString["state"] -eq $state) {
                $reader = [IO.StreamReader]::new($req.InputStream, $req.ContentEncoding)
                $jwt = $reader.ReadToEnd().Trim()
                $res.StatusCode = 204
                $res.Close()
            }
            else {
                $res.StatusCode = 404
                $res.Close()
            }
        }
    }
    finally { $listener.Stop() }

    if (-not $jwt) { throw "Did not receive a session token." }
    return $jwt
}

function Get-OrCreateVk {
    param([string]$Jwt)
    $email = Get-JwtEmail $Jwt
    $name  = "cowork-$email"

    # Reuse an existing key if the API returns its value; otherwise create.
    try {
        $existing = Invoke-BifrostApi -Path "/api/governance/virtual-keys" -Jwt $Jwt
        foreach ($vk in $existing.virtual_keys) {
            if ($vk.name -eq $name -and $vk.value) { return $vk.value }
        }
    } catch { }  # listing may be restricted; fall through to create

    $body = @{
        name             = $name
        provider_configs = @(@{
            provider       = "bedrock"
            allowed_models = $AllowedModels
            key_ids        = @("*")
        })
        budget = @{ max_limit = $MonthlyBudget; reset_duration = "1mo" }
    }
    $created = Invoke-BifrostApi -Path "/api/governance/virtual-keys" -Jwt $Jwt -Method POST -Body $body
    return $created.virtual_key.value
}

function Protect-VkFile {
    # Restrict the key file to the current user (chmod 600 equivalent).
    $acl = New-Object System.Security.AccessControl.FileSecurity
    $acl.SetAccessRuleProtection($true, $false)   # drop inheritance
    $rule = New-Object System.Security.AccessControl.FileSystemAccessRule(
        [Security.Principal.WindowsIdentity]::GetCurrent().Name,
        "FullControl", "Allow")
    $acl.AddAccessRule($rule)
    Set-Acl -Path $VkFile -AclObject $acl
}

# ---------------- Main ----------------
New-Item -ItemType Directory -Path $ConfDir -Force | Out-Null

if (-not (Test-HaveVk)) {
    $jwt = Get-SessionJwt
    $vk  = Get-OrCreateVk -Jwt $jwt
    Set-Content -Path $VkFile -Value $vk -NoNewline
    Protect-VkFile
    Write-Host "Virtual key provisioned." -ForegroundColor Green
}

$vk = (Get-Content $VkFile -Raw).Trim()

# On 401/403 from Bifrost, delete the cached key and re-run:
#   Remove-Item "$env:USERPROFILE\.config\cowork-bifrost\vk"; . .\cowork-login.ps1

$env:ANTHROPIC_BASE_URL  = $Bifrost
$env:ANTHROPIC_AUTH_TOKEN = $vk
Write-Host "ANTHROPIC_BASE_URL and ANTHROPIC_AUTH_TOKEN set for this session." -ForegroundColor Green
<#
.SYNOPSIS
  bifrost-session-login: obtain the user's Bifrost session/IdP token (via the
  Bifrost dashboard's own Okta login) and export it for Claude Cowork.

  No Okta admin access needed — this reuses Bifrost's existing Okta app.
  Flow: open browser -> Bifrost redirects to Okta -> user signs in ->
  callback page reads the Bifrost session token and posts it to this script.

.USAGE
  . .\bifrost-session-login.ps1     # dot-source so env vars persist
  # then launch Claude Cowork from the same shell
#>

$ErrorActionPreference = "Stop"

# ---------------- Config ----------------
$Bifrost = "https://bifrost.yourdomain.com"

# VERIFY THIS against your deployment: open the Bifrost dashboard, sign in,
# open DevTools (F12) -> Network tab, and find which request/endpoint the UI
# uses to get its bearer token (and which JSON field holds it). Adjust below.
$SessionEndpoint = "$Bifrost/api/auth/session"
$TokenField      = "token"

$ConfDir   = Join-Path $env:USERPROFILE ".config\cowork-bifrost"
$TokenFile = Join-Path $ConfDir "session-token"
$Port      = 8899
# ----------------------------------------

function Get-CachedToken {
    if (-not (Test-Path $TokenFile)) { return $null }
    $tok = (Get-Content $TokenFile -Raw).Trim()
    if (-not $tok) { return $null }

    # If it's a JWT, check expiry locally (with 5 min slack)
    $parts = $tok.Split(".")
    if ($parts.Count -eq 3) {
        try {
            $p = $parts[1].Replace('-', '+').Replace('_', '/')
            switch ($p.Length % 4) { 2 { $p += "==" } 3 { $p += "=" } }
            $claims = [Text.Encoding]::UTF8.GetString(
                        [Convert]::FromBase64String($p)) | ConvertFrom-Json
            if ($claims.exp) {
                $exp = [DateTimeOffset]::FromUnixTimeSeconds($claims.exp).UtcDateTime
                if ($exp -lt (Get-Date).ToUniversalTime().AddMinutes(5)) { return $null }
            }
        } catch { }  # opaque token — validate by probing instead
    }

    # Probe Bifrost to confirm the token is still accepted
    try {
        Invoke-RestMethod -Uri "$Bifrost/api/governance/virtual-keys" `
            -Headers @{ Authorization = "Bearer $tok" } -Method GET | Out-Null
        return $tok
    } catch {
        return $null   # 401/403 -> force re-login
    }
}

function Get-SessionToken {
    $state    = [Guid]::NewGuid().ToString("N")
    $listener = [System.Net.HttpListener]::new()
    $listener.Prefixes.Add("http://127.0.0.1:$Port/")
    $listener.Start()

    # Bifrost's /login will bounce through Okta if there's no active session.
    Start-Process "$Bifrost/login?redirect=http://localhost:$Port/callback"
    Write-Host "Complete the Okta sign-in in your browser..." -ForegroundColor Cyan

    $callbackHtml = @"
<!doctype html><html><body><h3>Finishing sign-in...</h3>
<script>
  fetch('$SessionEndpoint', {credentials:'include'})
    .then(r => { if (!r.ok) throw new Error('session fetch ' + r.status); return r.json(); })
    .then(d => fetch('/token?state=$state', {method:'POST', body:d.$TokenField}))
    .then(() => { document.body.innerHTML='<h3>Done — you can close this tab.</h3>'; window.close(); })
    .catch(e => { document.body.innerHTML =
      '<h3>Could not read the session token automatically ('+e+').</h3>' +
      '<p>Fallback: open DevTools on the Bifrost dashboard, copy your bearer token, ' +
      'and save it to %USERPROFILE%\\.config\\cowork-bifrost\\session-token</p>'; });
</script></body></html>
"@

    $token    = $null
    $deadline = (Get-Date).AddMinutes(5)
    try {
        while (-not $token -and (Get-Date) -lt $deadline) {
            $ctxTask = $listener.GetContextAsync()
            while (-not $ctxTask.AsyncWaitHandle.WaitOne(500)) {
                if ((Get-Date) -ge $deadline) { throw "Timed out waiting for login." }
            }
            $ctx = $ctxTask.Result
            $req = $ctx.Request; $res = $ctx.Response

            if ($req.HttpMethod -eq "GET" -and $req.Url.AbsolutePath -eq "/callback") {
                $b = [Text.Encoding]::UTF8.GetBytes($callbackHtml)
                $res.ContentType = "text/html"
                $res.OutputStream.Write($b, 0, $b.Length); $res.Close()
            }
            elseif ($req.HttpMethod -eq "POST" -and
                    $req.Url.AbsolutePath -eq "/token" -and
                    $req.QueryString["state"] -eq $state) {
                $rd = [IO.StreamReader]::new($req.InputStream, $req.ContentEncoding)
                $token = $rd.ReadToEnd().Trim()
                $res.StatusCode = 204; $res.Close()
            }
            else { $res.StatusCode = 404; $res.Close() }
        }
    }
    finally { $listener.Stop() }

    if (-not $token) { throw "Did not receive a session token." }
    return $token
}

function Protect-TokenFile {
    $acl = New-Object System.Security.AccessControl.FileSecurity
    $acl.SetAccessRuleProtection($true, $false)
    $rule = New-Object System.Security.AccessControl.FileSystemAccessRule(
        [Security.Principal.WindowsIdentity]::GetCurrent().Name, "FullControl", "Allow")
    $acl.AddAccessRule($rule)
    Set-Acl -Path $TokenFile -AclObject $acl
}

# ---------------- Main ----------------
New-Item -ItemType Directory -Path $ConfDir -Force | Out-Null

$tok = Get-CachedToken
if (-not $tok) {
    $tok = Get-SessionToken
    Set-Content -Path $TokenFile -Value $tok -NoNewline
    Protect-TokenFile
    Write-Host "Session token captured." -ForegroundColor Green
} else {
    Write-Host "Using cached session token." -ForegroundColor Green
}

$env:ANTHROPIC_BASE_URL   = $Bifrost
$env:ANTHROPIC_AUTH_TOKEN = $tok
Write-Host "ANTHROPIC_BASE_URL and ANTHROPIC_AUTH_TOKEN set. Launch Cowork from this shell." -ForegroundColor Green
$callbackHtml = @"
<!doctype html><html><body><h3>Checking Bifrost session...</h3>
<script>
  const already = new URLSearchParams(location.search).has('back');
  fetch('$SessionEndpoint', {credentials:'include'})
    .then(r => {
      if (r.status === 401 || r.status === 403) {
        if (already) throw new Error('still unauthenticated after login');
        // Not signed in: go log in, then come back here
        location.href = '$Bifrost/login?redirect=' +
          encodeURIComponent('http://localhost:$Port/callback?back=1');
        return new Promise(()=>{});  // stop the chain during redirect
      }
      if (!r.ok) throw new Error('session fetch ' + r.status);
      return r.json();
    })
    .then(d => fetch('/token?state=$state', {method:'POST', body:d.$TokenField}))
    .then(() => { document.body.innerHTML='<h3>Done — you can close this tab.</h3>'; window.close(); })
    .catch(e => { document.body.innerHTML =
      '<h3>Could not read the session token ('+e+').</h3>' +
      '<p>Fallback: copy your bearer token from DevTools on the Bifrost dashboard and save it to ' +
      '%USERPROFILE%\\.config\\cowork-bifrost\\session-token</p>'; });
</script></body></html>
"@
Scroll to Top