1  import * as crypto from 'crypto'
  2  import * as http from 'http'
  3  import { IncomingMessage, ServerResponse } from 'http'
  4  import * as url from 'url'
  5  
  6  import { OAUTH_CONFIG } from '../constants/oauth.js'
  7  import { openBrowser } from '../utils/browser.js'
  8  import { logEvent } from '../services/statsig.js'
  9  import { logError } from '../utils/log.js'
 10  import { resetAnthropicClient } from './claude.js'
 11  import {
 12    AccountInfo,
 13    getGlobalConfig,
 14    saveGlobalConfig,
 15    normalizeApiKeyForConfig,
 16  } from '../utils/config.js'
 17  
 18  // Base64URL encoding function (RFC 4648)
 19  function base64URLEncode(buffer: Buffer): string {
 20    return buffer
 21      .toString('base64')
 22      .replace(/\+/g, '-')
 23      .replace(/\//g, '_')
 24      .replace(/=/g, '')
 25  }
 26  
 27  function generateCodeVerifier(): string {
 28    return base64URLEncode(crypto.randomBytes(32))
 29  }
 30  
 31  async function generateCodeChallenge(verifier: string): Promise<string> {
 32    const encoder = new TextEncoder()
 33    const data = encoder.encode(verifier)
 34    const digest = await crypto.subtle.digest('SHA-256', data)
 35    return base64URLEncode(Buffer.from(digest))
 36  }
 37  
 38  type OAuthTokenExchangeResponse = {
 39    access_token: string
 40    account?: {
 41      uuid: string
 42      email_address: string
 43    }
 44    organization?: {
 45      uuid: string
 46      name: string
 47    }
 48  }
 49  
 50  export type OAuthResult = {
 51    accessToken: string
 52  }
 53  
 54  export class OAuthService {
 55    private server: http.Server | null = null
 56    private codeVerifier: string
 57    private expectedState: string | null = null
 58    private pendingCodePromise: {
 59      resolve: (result: {
 60        authorizationCode: string
 61        useManualRedirect: boolean
 62      }) => void
 63      reject: (err: Error) => void
 64    } | null = null
 65  
 66    constructor() {
 67      this.codeVerifier = generateCodeVerifier()
 68    }
 69  
 70    private generateAuthUrls(
 71      codeChallenge: string,
 72      state: string,
 73    ): { autoUrl: string; manualUrl: string } {
 74      function makeUrl(isManual: boolean): string {
 75        const authUrl = new URL(OAUTH_CONFIG.AUTHORIZE_URL)
 76        authUrl.searchParams.append('client_id', OAUTH_CONFIG.CLIENT_ID)
 77        authUrl.searchParams.append('response_type', 'code')
 78        authUrl.searchParams.append(
 79          'redirect_uri',
 80          isManual
 81            ? OAUTH_CONFIG.MANUAL_REDIRECT_URL
 82            : `http://localhost:${OAUTH_CONFIG.REDIRECT_PORT}/callback`,
 83        )
 84        authUrl.searchParams.append('scope', OAUTH_CONFIG.SCOPES.join(' '))
 85        authUrl.searchParams.append('code_challenge', codeChallenge)
 86        authUrl.searchParams.append('code_challenge_method', 'S256')
 87        authUrl.searchParams.append('state', state)
 88        return authUrl.toString()
 89      }
 90  
 91      return {
 92        autoUrl: makeUrl(false),
 93        manualUrl: makeUrl(true),
 94      }
 95    }
 96  
 97    async startOAuthFlow(
 98      authURLHandler: (url: string) => Promise<void>,
 99    ): Promise<OAuthResult> {
100      const codeChallenge = await generateCodeChallenge(this.codeVerifier)
101      const state = base64URLEncode(crypto.randomBytes(32))
102      this.expectedState = state
103      const { autoUrl, manualUrl } = this.generateAuthUrls(codeChallenge, state)
104  
105      const onReady = async () => {
106        await authURLHandler(manualUrl)
107        await openBrowser(autoUrl)
108      }
109  
110      const { authorizationCode, useManualRedirect } = await new Promise<{
111        authorizationCode: string
112        useManualRedirect: boolean
113      }>((resolve, reject) => {
114        this.pendingCodePromise = { resolve, reject }
115        this.startLocalServer(state, onReady)
116      })
117  
118      // Exchange code for tokens
119      const {
120        access_token: accessToken,
121        account,
122        organization,
123      } = await this.exchangeCodeForTokens(
124        authorizationCode,
125        state,
126        useManualRedirect,
127      )
128  
129      // Store account info
130      if (account) {
131        const accountInfo: AccountInfo = {
132          accountUuid: account.uuid,
133          emailAddress: account.email_address,
134          organizationUuid: organization?.uuid,
135        }
136        const config = getGlobalConfig()
137        config.oauthAccount = accountInfo
138        saveGlobalConfig(config)
139      }
140  
141      return { accessToken }
142    }
143  
144    private startLocalServer(state: string, onReady?: () => void): void {
145      if (this.server) {
146        this.closeServer()
147      }
148      this.server = http.createServer(
149        (req: IncomingMessage, res: ServerResponse) => {
150          const parsedUrl = url.parse(req.url || '', true)
151  
152          if (parsedUrl.pathname === '/callback') {
153            const authorizationCode = parsedUrl.query.code as string
154            const returnedState = parsedUrl.query.state as string
155  
156            if (!authorizationCode) {
157              res.writeHead(400)
158              res.end('Authorization code not found')
159              if (this.pendingCodePromise) {
160                this.pendingCodePromise.reject(
161                  new Error('No authorization code received'),
162                )
163              }
164              return
165            }
166  
167            if (returnedState !== state) {
168              res.writeHead(400)
169              res.end('Invalid state parameter')
170              if (this.pendingCodePromise) {
171                this.pendingCodePromise.reject(
172                  new Error('Invalid state parameter'), // Possible CSRF attack
173                )
174              }
175              return
176            }
177  
178            res.writeHead(302, {
179              Location: OAUTH_CONFIG.SUCCESS_URL,
180            })
181            res.end()
182  
183            // Track which path the user is taking (automatic browser redirect)
184            logEvent('tengu_oauth_automatic_redirect', {})
185  
186            this.processCallback({
187              authorizationCode,
188              state,
189              useManualRedirect: false,
190            })
191          } else {
192            res.writeHead(404)
193            res.end()
194          }
195        },
196      )
197  
198      this.server.listen(OAUTH_CONFIG.REDIRECT_PORT, async () => {
199        onReady?.()
200      })
201  
202      this.server.on('error', (err: Error) => {
203        const portError = err as NodeJS.ErrnoException
204        if (portError.code === 'EADDRINUSE') {
205          const error = new Error(
206            `Port ${OAUTH_CONFIG.REDIRECT_PORT} is already in use. Please ensure no other applications are using this port.`,
207          )
208          logError(error)
209          this.closeServer()
210          if (this.pendingCodePromise) {
211            this.pendingCodePromise.reject(error)
212          }
213          return
214        } else {
215          logError(err)
216          this.closeServer()
217          if (this.pendingCodePromise) {
218            this.pendingCodePromise.reject(err)
219          }
220          return
221        }
222      })
223    }
224  
225    private async exchangeCodeForTokens(
226      authorizationCode: string,
227      state: string,
228      useManualRedirect: boolean = false,
229    ): Promise<OAuthTokenExchangeResponse> {
230      const requestBody = {
231        grant_type: 'authorization_code',
232        code: authorizationCode,
233        redirect_uri: useManualRedirect
234          ? OAUTH_CONFIG.MANUAL_REDIRECT_URL
235          : `http://localhost:${OAUTH_CONFIG.REDIRECT_PORT}/callback`,
236        client_id: OAUTH_CONFIG.CLIENT_ID,
237        code_verifier: this.codeVerifier,
238        state,
239      }
240  
241      const response = await fetch(OAUTH_CONFIG.TOKEN_URL, {
242        method: 'POST',
243        headers: {
244          'Content-Type': 'application/json',
245        },
246        body: JSON.stringify(requestBody),
247      })
248  
249      if (!response.ok) {
250        throw new Error(`Token exchange failed: ${response.statusText}`)
251      }
252  
253      const data = await response.json()
254      return data
255    }
256  
257    processCallback({
258      authorizationCode,
259      state,
260      useManualRedirect,
261    }: {
262      authorizationCode: string
263      state: string
264      useManualRedirect: boolean
265    }): void {
266      this.closeServer()
267  
268      if (state !== this.expectedState) {
269        if (this.pendingCodePromise) {
270          this.pendingCodePromise.reject(
271            new Error('Invalid state parameter'), // Possible CSRF attack
272          )
273          this.pendingCodePromise = null
274        }
275        return
276      }
277  
278      if (this.pendingCodePromise) {
279        this.pendingCodePromise.resolve({ authorizationCode, useManualRedirect })
280        this.pendingCodePromise = null
281      }
282    }
283  
284    private closeServer(): void {
285      if (this.server) {
286        this.server.close()
287        this.server = null
288      }
289    }
290  }
291  
292  export async function createAndStoreApiKey(
293    accessToken: string,
294  ): Promise<string | null> {
295    try {
296      // Call create_api_key endpoint
297      const createApiKeyResp = await fetch(OAUTH_CONFIG.API_KEY_URL, {
298        method: 'POST',
299        headers: { Authorization: `Bearer ${accessToken}` },
300      })
301  
302      let apiKeyData
303      let errorText = ''
304  
305      try {
306        apiKeyData = await createApiKeyResp.json()
307      } catch (_e) {
308        // If response is not valid JSON, get as text for error logging
309        errorText = await createApiKeyResp.text()
310      }
311  
312      logEvent('tengu_oauth_api_key', {
313        status: createApiKeyResp.ok ? 'success' : 'failure',
314        statusCode: createApiKeyResp.status.toString(),
315        error: createApiKeyResp.ok ? '' : errorText || JSON.stringify(apiKeyData),
316      })
317  
318      if (createApiKeyResp.ok && apiKeyData && apiKeyData.raw_key) {
319        const apiKey = apiKeyData.raw_key
320  
321        // Store in global config
322        const config = getGlobalConfig()
323  
324        // Store as primary API key
325        config.primaryApiKey = apiKey
326  
327        // Add to approved list
328        if (!config.customApiKeyResponses) {
329          config.customApiKeyResponses = { approved: [], rejected: [] }
330        }
331        if (!config.customApiKeyResponses.approved) {
332          config.customApiKeyResponses.approved = []
333        }
334  
335        const normalizedKey = normalizeApiKeyForConfig(apiKey)
336        if (!config.customApiKeyResponses.approved.includes(normalizedKey)) {
337          config.customApiKeyResponses.approved.push(normalizedKey)
338        }
339  
340        // Save config
341        saveGlobalConfig(config)
342  
343        // Reset the Anthropic client to force creation with new API key
344        resetAnthropicClient()
345  
346        return apiKey
347      }
348  
349      return null
350    } catch (error) {
351      logEvent('tengu_oauth_api_key', {
352        status: 'failure',
353        statusCode: 'exception',
354        error: error instanceof Error ? error.message : String(error),
355      })
356      throw error
357    }
358  }