@@ -18,6 +18,7 @@ type OrgInfo = {
1818 name : string
1919 image : string | null
2020 plan : 'opensource' | 'team' | 'enterprise'
21+ slug : string
2122}
2223
2324type OrganizationsRecord = {
@@ -46,6 +47,16 @@ async function getOrganizations(apiKey: string): Promise<OrganizationsRecord | n
4647 return orgs
4748}
4849
50+ const orgSlugByApiKey = new Map < string , string > ( )
51+
52+ function getDefaultOrg ( organizations : OrganizationsRecord ) : OrgInfo | null {
53+ const org = Object . values ( organizations . organizations ) [ 0 ]
54+ if ( ! org || ! org . slug ) {
55+ return null
56+ }
57+ return org
58+ }
59+
4960export async function activate ( context : vscode . ExtensionContext , disposables : Array < vscode . Disposable > ) {
5061 //#region file path/watching
5162 // responsible for watching files to know when to sync from disk
@@ -120,12 +131,13 @@ export async function activate(context: vscode.ExtensionContext, disposables: Ar
120131 const sessionOnDisk : typeof liveSessions = new Map < vscode . AuthenticationSession [ 'accessToken' ] , vscode . AuthenticationSession > ( )
121132 if ( typeof apiKey === 'string' && apiKey . length > 0 && apiKey !== SOCKET_PUBLIC_API_TOKEN ) {
122133 const organizations = await getOrganizations ( apiKey )
123- const org = Object . values ( organizations ! . organizations ) [ 0 ]
124- if ( org ) {
134+ const defaultOrg = organizations ? getDefaultOrg ( organizations ) : null
135+ if ( defaultOrg ) {
125136 sessionOnDisk . set (
126137 apiKey ,
127- sessionFromAPIKey ( apiKey , org )
138+ sessionFromAPIKey ( apiKey , defaultOrg )
128139 )
140+ orgSlugByApiKey . set ( apiKey , defaultOrg . slug )
129141 }
130142 }
131143 let added : Array < vscode . AuthenticationSession > = [ ]
@@ -176,23 +188,41 @@ export async function activate(context: vscode.ExtensionContext, disposables: Ar
176188 return Array . from ( liveSessions . values ( ) )
177189 } ,
178190 async createSession ( scopes : readonly string [ ] , options : vscode . AuthenticationProviderSessionOptions ) : Promise < vscode . AuthenticationSession > {
179- let organizations : OrganizationsRecord
191+ let organizations : OrganizationsRecord | null = null
192+ let defaultOrg : OrgInfo | null = null
180193 let apiKey : string = await vscode . window . showInputBox ( {
181194 title : 'Socket Security API Token' ,
182195 placeHolder : 'Leave this blank to stay logged out' ,
183196 ignoreFocusOut : true ,
184197 prompt : 'Enter your API token from https://socket.dev/' ,
185198 async validateInput ( value ) {
186- if ( ! value ) return
187- organizations = ( await getOrganizations ( value ) ) !
188- if ( ! organizations ) return 'Invalid API key'
199+ if ( ! value ) {
200+ return
201+ }
202+
203+ organizations = await getOrganizations ( value )
204+ if ( ! organizations ) {
205+ return 'Invalid API key'
206+ }
207+
208+ defaultOrg = getDefaultOrg ( organizations )
209+ if ( ! defaultOrg ) {
210+ return 'No organizations found for API key'
211+ }
189212 }
190213 } ) ?? ''
191214 if ( ! apiKey ) {
192215 throw new Error ( 'User did not want to provide an API key' )
193216 }
194- const org = Object . values ( organizations ! . organizations ) [ 0 ]
195- const session = sessionFromAPIKey ( apiKey , org )
217+ if ( ! organizations ) {
218+ organizations = await getOrganizations ( apiKey )
219+ }
220+ defaultOrg = defaultOrg ?? ( organizations ? getDefaultOrg ( organizations ) : null )
221+ if ( ! defaultOrg ) {
222+ throw new Error ( 'No organizations found for API key' )
223+ }
224+ const session = sessionFromAPIKey ( apiKey , defaultOrg )
225+ orgSlugByApiKey . set ( apiKey , defaultOrg . slug )
196226 let oldSessions = Array . from ( liveSessions . values ( ) )
197227 await syncLiveSessionToDisk ( session )
198228 liveSessions = new Map ( [
@@ -228,6 +258,7 @@ export async function activate(context: vscode.ExtensionContext, disposables: Ar
228258 }
229259 } catch { }
230260 if ( session ) {
261+ orgSlugByApiKey . delete ( session . accessToken )
231262 diskSessionsChanges . fire ( {
232263 added : [ ] ,
233264 changed : [ ] ,
@@ -271,6 +302,24 @@ export async function getAPIKey() {
271302 }
272303}
273304
305+ export async function getOrgSlug ( apiKey ?: string ) {
306+ const resolvedApiKey = apiKey ?? await getAPIKey ( )
307+ if ( ! resolvedApiKey || resolvedApiKey === SOCKET_PUBLIC_API_TOKEN ) {
308+ return null
309+ }
310+ const cached = orgSlugByApiKey . get ( resolvedApiKey )
311+ if ( cached ) {
312+ return cached
313+ }
314+ const organizations = await getOrganizations ( resolvedApiKey )
315+ const defaultOrg = organizations ? getDefaultOrg ( organizations ) : null
316+ if ( ! defaultOrg ) {
317+ return null
318+ }
319+ orgSlugByApiKey . set ( resolvedApiKey , defaultOrg . slug )
320+ return defaultOrg . slug
321+ }
322+
274323export function getAuthHeader ( apiKey : string ) {
275324 return `Bearer ${ apiKey } `
276325}
@@ -289,4 +338,3 @@ function sessionFromAPIKey(apiKey: string, org: OrgInfo) {
289338 scopes : [ ] ,
290339 }
291340}
292-
0 commit comments