diff --git a/src/Commands/PollingDiscoveryStatus.ts b/src/Commands/PollingDiscoveryStatus.ts new file mode 100644 index 00000000..7265d526 --- /dev/null +++ b/src/Commands/PollingDiscoveryStatus.ts @@ -0,0 +1,83 @@ +import { RestDiscoveryOptions } from 'src/Discovery'; +import { DiscoveryPollingFactory } from 'src/Discovery/DiscoveryPollingFactory'; +import { ErrorMessageFactory, logger } from 'src/Utils'; +import { container } from 'tsyringe'; +import { Arguments, Argv, CommandModule } from 'yargs'; + +export class PollingDiscoveryStatus implements CommandModule { + public readonly command = 'discovery:polling [options] '; + public readonly describe = + 'Allows to configure a polling of discovery status.'; + + public builder(argv: Argv): Argv { + return argv + .option('token', { + alias: 't', + describe: 'Bright API-key', + requiresArg: true, + demandOption: true + }) + .option('project', { + alias: 'p', + describe: 'ID of the project', + string: true, + requiresArg: true, + demandOption: true + }) + .option('interval', { + requiresArg: true, + describe: + 'The sampling interval between status checks. ' + + 'Eg: 60, "2min", "10h", "7d". A numeric value is interpreted as a milliseconds count.', + default: 5000 + }) + .option('timeout', { + requiresArg: true, + describe: + 'Period of time between the end of a timeout period or completion of a discovery status request, and the next request for status. ' + + 'Eg: 60, "2min", "10h", "7d". A numeric value is interpreted as a milliseconds count.' + }) + .positional('discoveryId', { + describe: 'ID of an existing discovery.', + demandOption: true, + type: 'string' + }) + .middleware((args: Arguments) => + container.register(RestDiscoveryOptions, { + useValue: { + insecure: args.insecure as boolean, + baseURL: args.api as string, + apiKey: args.token as string, + proxyURL: (args.proxyBright ?? args.proxy) as string, + timeout: args.timeout as number + } + }) + ); + } + + public async handler(args: Arguments): Promise { + try { + const pollingFactory = container.resolve( + DiscoveryPollingFactory + ); + const polling = pollingFactory.create({ + discoveryId: args.discoveryId as string, + projectId: args.project as string, + timeout: args.timeout as number, + interval: args.interval as number + }); + + await polling.start(); + + process.exit(0); + } catch (error) { + logger.error( + ErrorMessageFactory.genericCommandError({ + error, + command: 'discovery:polling' + }) + ); + process.exit(1); + } + } +} diff --git a/src/Discovery/DefaultDiscoveryPollingFactory.ts b/src/Discovery/DefaultDiscoveryPollingFactory.ts new file mode 100644 index 00000000..66b952eb --- /dev/null +++ b/src/Discovery/DefaultDiscoveryPollingFactory.ts @@ -0,0 +1,20 @@ +import { + DiscoveryPollingConfig, + DiscoveryPollingFactory +} from './DiscoveryPollingFactory'; +import { Polling } from '../Utils/Polling'; +import { DiscoveryPolling } from './DiscoveryPolling'; +import { Discoveries } from './Discoveries'; +import { inject, injectable } from 'tsyringe'; + +@injectable() +export class DefaultDiscoveryPollingFactory implements DiscoveryPollingFactory { + constructor( + @inject(Discoveries) + private readonly discoveries: Discoveries + ) {} + + public create(options: DiscoveryPollingConfig): Polling { + return new DiscoveryPolling(options, this.discoveries); + } +} diff --git a/src/Discovery/Discoveries.ts b/src/Discovery/Discoveries.ts index b491570b..96e4ae3c 100644 --- a/src/Discovery/Discoveries.ts +++ b/src/Discovery/Discoveries.ts @@ -1,3 +1,5 @@ +import { DiscoveryView } from './DiscoveryView'; + export interface DiscoveryConfig { name: string; authObjectId?: string; @@ -34,6 +36,12 @@ export interface Discoveries { stop(projectId: string, discoveryId: string): Promise; delete(projectId: string, discoveryId: string): Promise; + + get( + projectId: string, + discoveryId: string, + options?: { signal?: AbortSignal } + ): Promise; } export const Discoveries: unique symbol = Symbol('Discoveries'); diff --git a/src/Discovery/DiscoveryPolling.spec.ts b/src/Discovery/DiscoveryPolling.spec.ts new file mode 100644 index 00000000..438e8f92 --- /dev/null +++ b/src/Discovery/DiscoveryPolling.spec.ts @@ -0,0 +1,169 @@ +import 'reflect-metadata'; +import { Logger, logger } from '../Utils'; +import { DiscoveryPollingConfig } from './DiscoveryPollingFactory'; +import { DiscoveryView, DiscoveryStatus } from './DiscoveryView'; +import { Discoveries } from './Discoveries'; +import { DiscoveryPolling } from './DiscoveryPolling'; +import { instance, mock, reset, spy, verify, when } from 'ts-mockito'; +import { setTimeout } from 'node:timers/promises'; + +describe('DiscoveryPolling', () => { + const discoveryId = 'hAXZjjahZqpvgK3yNEdp6t'; + const projectId = 'hADZjiahZqpvgK3yNEdp8b'; + + const firstResponse: DiscoveryView = { + id: discoveryId, + name: 'some name', + status: DiscoveryStatus.RUNNING + }; + + const discoveryManagerMock = mock(); + + let loggerSpy!: Logger; + + beforeEach(() => { + loggerSpy = spy(logger); + }); + + afterEach(() => { + reset(discoveryManagerMock, loggerSpy); + }); + + describe('constructor', () => { + it('should warn if timeout is not specified', () => { + // arrange + const options: DiscoveryPollingConfig = { + projectId, + discoveryId + }; + + // act + new DiscoveryPolling(options, instance(discoveryManagerMock)); + + // assert + verify( + loggerSpy.warn( + `Warning: It looks like you've been running polling without "timeout" option.` + ) + ).once(); + verify( + loggerSpy.warn( + `The recommended way to install polling with a minimal timeout: 10-20min.` + ) + ).once(); + }); + + it('should warn if interval is less than 10s', () => { + // arrange + const options = { + discoveryId, + projectId, + interval: 5000 + }; + + // act + new DiscoveryPolling(options, instance(discoveryManagerMock)); + + // assert + verify( + loggerSpy.warn( + `Warning: The minimal value for polling interval is 10 seconds.` + ) + ).once(); + }); + }); + + describe('start', () => { + const options: DiscoveryPollingConfig = { + discoveryId, + projectId, + interval: 1 + }; + const spiedOptions = spy(options); + + let sut!: DiscoveryPolling; + + beforeEach(() => { + sut = new DiscoveryPolling(options, instance(discoveryManagerMock)); + }); + + afterEach(() => reset(spiedOptions)); + + it.each([ + DiscoveryStatus.DONE, + DiscoveryStatus.DISRUPTED, + DiscoveryStatus.FAILED, + DiscoveryStatus.STOPPED + ])( + 'should start polling and stop on discovery status changed to "%s"', + async (status) => { + // arrange + when(discoveryManagerMock.get(projectId, discoveryId)) + .thenResolve(firstResponse) + .thenResolve({ ...firstResponse, status }); + + // act + await sut.start(); + + // assert + verify(discoveryManagerMock.get(projectId, discoveryId)).twice(); + verify( + loggerSpy.log( + `The discovery has been finished with status: ${status}.` + ) + ).once(); + } + ); + + it('should start polling and stop on timeout', async () => { + // arrange + const timeout = 1500; + const interval = 1000; + when(spiedOptions.timeout).thenReturn(timeout); + when(spiedOptions.interval).thenReturn(interval); + + when(discoveryManagerMock.get(projectId, discoveryId)).thenResolve( + firstResponse + ); + + // act + jest.useFakeTimers(); + const promise = sut.start(); + await setTimeout(10); + jest.runAllTimers(); + await promise; + jest.useRealTimers(); + + // assert + verify(discoveryManagerMock.get(projectId, discoveryId)).once(); + verify(loggerSpy.log('Polling has been stopped by timeout.')).once(); + }); + }); + + describe('stop', () => { + it('should stop polling', async () => { + // arrange + const sut = new DiscoveryPolling( + { + projectId, + discoveryId, + interval: 1000 + }, + instance(discoveryManagerMock) + ); + + when(discoveryManagerMock.get(projectId, discoveryId)).thenResolve( + firstResponse + ); + + // act + const start = sut.start(); + await setTimeout(10); + await sut.stop(); + await start; + + // assert + verify(discoveryManagerMock.get(projectId, discoveryId)).once(); + }); + }); +}); diff --git a/src/Discovery/DiscoveryPolling.ts b/src/Discovery/DiscoveryPolling.ts new file mode 100644 index 00000000..2e0e058f --- /dev/null +++ b/src/Discovery/DiscoveryPolling.ts @@ -0,0 +1,172 @@ +import { Polling } from '../Utils/Polling'; +import { Backoff, ErrorMessageFactory, logger } from '../Utils'; +import { DiscoveryPollingConfig } from './DiscoveryPollingFactory'; +import { Discoveries } from './Discoveries'; +import { DiscoveryStatus, DiscoveryView } from './DiscoveryView'; +import axios from 'axios'; +import { setTimeout as asyncSetTimeout } from 'node:timers/promises'; + +export class DiscoveryPolling implements Polling { + private timeoutDescriptor?: NodeJS.Timeout; + private defaultInterval: number = 10000; + private readonly DEFAULT_RECONNECT_TIMES = 20; + private abortController = new AbortController(); + + constructor( + private readonly options: DiscoveryPollingConfig, + private readonly discoveryManager: Discoveries + ) { + if (!this.options.timeout) { + logger.warn( + `Warning: It looks like you've been running polling without "timeout" option.` + ); + logger.warn( + `The recommended way to install polling with a minimal timeout: 10-20min.` + ); + } + + if (this.options.interval) { + if (this.options.interval < this.defaultInterval) { + logger.warn( + `Warning: The minimal value for polling interval is 10 seconds.` + ); + } + } + } + + public async start(): Promise { + try { + logger.log('Starting polling...'); + this.initializePolling(); + await this.runPollingLoop(); + } catch (error) { + this.handleError(error); + } finally { + await this.stop(); + } + } + + // eslint-disable-next-line @typescript-eslint/require-await + public async stop(): Promise { + this.abortController.abort(); + clearTimeout(this.timeoutDescriptor); + } + + private initializePolling(): void { + if (this.options.timeout) { + this.setTimeout(); + } + } + + private async runPollingLoop(): Promise { + for await (const discovery of this.poll()) { + const shouldContinue = await this.processDiscoveryView(discovery); + if (!shouldContinue) break; + } + } + + private handleError(error: unknown): void { + if (!this.abortController.signal.aborted) { + ErrorMessageFactory.genericCommandError({ + error, + command: 'discovery:polling' + }); + process.exit(1); + } + } + + private setTimeout(timeout: number = this.options.timeout): void { + this.timeoutDescriptor = setTimeout(() => { + this.abortController.abort(); + logger.log('Polling has been stopped by timeout.'); + }, timeout); + logger.debug(`The polling timeout has been set to %d ms.`, timeout); + } + + private async *poll(): AsyncIterableIterator { + while (!this.abortController.signal.aborted) { + const backoff = this.createBackoff(); + + const view: DiscoveryView = await backoff.execute(() => + this.discoveryManager.get( + this.options.projectId, + this.options.discoveryId + ) + ); + + yield view; + + await this.delay(); + } + } + + private isFinished(status: DiscoveryStatus): boolean { + return ( + status === DiscoveryStatus.DONE || + status === DiscoveryStatus.STOPPED || + status === DiscoveryStatus.DISRUPTED || + status === DiscoveryStatus.FAILED + ); + } + + private async delay(): Promise { + const interval = this.options.interval ?? this.defaultInterval; + await asyncSetTimeout(interval, false, { + signal: this.abortController.signal + }); + } + + private createBackoff(): Backoff { + return new Backoff( + this.DEFAULT_RECONNECT_TIMES, + (err: unknown) => + (axios.isAxiosError(err) && err.status > 500) || + [ + 'ECONNRESET', + 'ENETDOWN', + 'ENETUNREACH', + 'ETIMEDOUT', + 'ECONNREFUSED', + 'ENOTFOUND', + 'EAI_AGAIN', + 'ESOCKETTIMEDOUT' + ].includes((err as NodeJS.ErrnoException).code) + ); + } + + private handleDiscoveryStatus(status: DiscoveryStatus): void { + const statusMessages = { + [DiscoveryStatus.RUNNING]: 'Discovery is running.', + [DiscoveryStatus.PENDING]: 'Discovery is pending.', + [DiscoveryStatus.SCHEDULED]: 'Discovery is scheduled.', + [DiscoveryStatus.QUEUED]: 'Discovery is queued.', + [DiscoveryStatus.DONE]: 'Discovery has been completed.', + [DiscoveryStatus.STOPPED]: 'Discovery has been stopped.', + [DiscoveryStatus.DISRUPTED]: 'Discovery has been disrupted.', + [DiscoveryStatus.FAILED]: 'Discovery has failed.' + }; + + const message = statusMessages[status] || `Discovery status is ${status}.`; + logger.log(message); + } + + private processDiscoveryView(discovery: DiscoveryView | null): boolean { + if (!discovery) { + logger.log('The discovery has not been found.'); + + return false; + } + + this.handleDiscoveryStatus(discovery.status); + + if (this.isFinished(discovery.status)) { + logger.log( + `The discovery has been finished with status: ${discovery.status}.` + ); + + return false; + } + + return true; + } +} diff --git a/src/Discovery/DiscoveryPollingFactory.ts b/src/Discovery/DiscoveryPollingFactory.ts new file mode 100644 index 00000000..df73a02f --- /dev/null +++ b/src/Discovery/DiscoveryPollingFactory.ts @@ -0,0 +1,16 @@ +import { Polling } from '../Utils/Polling'; + +export interface DiscoveryPollingConfig { + timeout?: number; + interval?: number; + discoveryId: string; + projectId: string; +} + +export interface DiscoveryPollingFactory { + create(options: DiscoveryPollingConfig): Polling; +} + +export const DiscoveryPollingFactory: unique symbol = Symbol( + 'DiscoveryPollingFactory' +); diff --git a/src/Discovery/DiscoveryView.ts b/src/Discovery/DiscoveryView.ts new file mode 100644 index 00000000..d59b093d --- /dev/null +++ b/src/Discovery/DiscoveryView.ts @@ -0,0 +1,16 @@ +export interface DiscoveryView { + id: string; + name: string; + status: DiscoveryStatus; +} + +export enum DiscoveryStatus { + RUNNING = 'running', + PENDING = 'pending', + STOPPED = 'stopped', + FAILED = 'failed', + DONE = 'done', + DISRUPTED = 'disrupted', + SCHEDULED = 'scheduled', + QUEUED = 'queued' +} diff --git a/src/Discovery/RestDiscoveries.ts b/src/Discovery/RestDiscoveries.ts index 5ee1efb0..301ee49d 100644 --- a/src/Discovery/RestDiscoveries.ts +++ b/src/Discovery/RestDiscoveries.ts @@ -9,6 +9,7 @@ import { } from './Discoveries'; import { ProxyFactory } from '../Utils'; import { CliInfo } from '../Config'; +import { DiscoveryView } from './DiscoveryView'; import { delay, inject, injectable } from 'tsyringe'; import axios, { Axios } from 'axios'; import http from 'node:http'; @@ -93,6 +94,19 @@ export class RestDiscoveries implements Discoveries { ); } + public async get( + projectId: string, + discoveryId: string, + options?: { signal?: AbortSignal } + ): Promise { + const res = await this.client.get( + `/api/v2/projects/${projectId}/discoveries/${discoveryId}`, + { signal: options?.signal } + ); + + return res.data; + } + private async prepareConfig({ headers, ...rest }: DiscoveryConfig): Promise< Omit & { headers: Header[]; diff --git a/src/Discovery/index.ts b/src/Discovery/index.ts index 35cbb1f5..66da911d 100644 --- a/src/Discovery/index.ts +++ b/src/Discovery/index.ts @@ -1,2 +1,4 @@ export * from './Discoveries'; export * from './RestDiscoveries'; +export * from './DiscoveryPollingFactory'; +export * from './DefaultDiscoveryPollingFactory'; diff --git a/src/Scan/BasePolling.spec.ts b/src/Scan/BasePolling.spec.ts index eda37ec2..4a3da868 100644 --- a/src/Scan/BasePolling.spec.ts +++ b/src/Scan/BasePolling.spec.ts @@ -83,9 +83,10 @@ describe('BasePolling', () => { ); // assert - verify(loggerSpy.warn(`Warning: polling interval is too small.`)).once(); verify( - loggerSpy.warn(`The recommended way to set polling interval to 10s.`) + loggerSpy.warn( + `Warning: The minimal value for polling interval is 10 seconds.` + ) ).once(); }); }); diff --git a/src/Scan/BasePolling.ts b/src/Scan/BasePolling.ts index 5c24ac96..9664b60d 100644 --- a/src/Scan/BasePolling.ts +++ b/src/Scan/BasePolling.ts @@ -1,5 +1,5 @@ import { Scans, ScanState, ScanStatus } from './Scans'; -import { Polling } from './Polling'; +import { Polling } from '../Utils/Polling'; import { Breakpoint } from './Breakpoint'; import { Backoff, logger } from '../Utils'; import { PollingConfig } from './PollingFactory'; @@ -33,8 +33,9 @@ export class BasePolling implements Polling { if (this.options.interval) { if (this.options.interval < this.defaultInterval) { - logger.warn(`Warning: polling interval is too small.`); - logger.warn(`The recommended way to set polling interval to 10s.`); + logger.warn( + `Warning: The minimal value for polling interval is 10 seconds.` + ); } } diff --git a/src/Scan/DefaultPollingFactory.ts b/src/Scan/DefaultPollingFactory.ts index 933ff909..3654682c 100644 --- a/src/Scan/DefaultPollingFactory.ts +++ b/src/Scan/DefaultPollingFactory.ts @@ -1,6 +1,6 @@ import { PollingConfig, PollingFactory } from './PollingFactory'; import { Scans } from './Scans'; -import { Polling } from './Polling'; +import { Polling } from '../Utils/Polling'; import { BasePolling } from './BasePolling'; import { BreakpointFactory } from './BreakpointFactory'; import { inject, injectable } from 'tsyringe'; diff --git a/src/Scan/PollingFactory.ts b/src/Scan/PollingFactory.ts index 875161fb..cac42202 100644 --- a/src/Scan/PollingFactory.ts +++ b/src/Scan/PollingFactory.ts @@ -1,4 +1,4 @@ -import { Polling } from './Polling'; +import { Polling } from '../Utils/Polling'; import { BreakpointType } from './BreakpointType'; export interface PollingConfig { diff --git a/src/Scan/index.ts b/src/Scan/index.ts index af656d7d..ec31d002 100644 --- a/src/Scan/index.ts +++ b/src/Scan/index.ts @@ -2,7 +2,7 @@ export * from './BreakpointFactory'; export * from './DefaultBreakpointFactory'; export * from './PollingFactory'; export * from './DefaultPollingFactory'; -export * from './Polling'; +export * from '../Utils/Polling'; export * from './Scans'; export * from './BreakpointType'; export * from './Breakpoint'; diff --git a/src/Scan/Polling.ts b/src/Utils/Polling.ts similarity index 100% rename from src/Scan/Polling.ts rename to src/Utils/Polling.ts diff --git a/src/container.ts b/src/container.ts index b75a761a..4a93fd01 100644 --- a/src/container.ts +++ b/src/container.ts @@ -56,7 +56,12 @@ import { ServerRepeaterLauncher } from './Repeater'; import { ProxyFactory, DefaultProxyFactory } from './Utils'; -import { Discoveries, RestDiscoveries } from './Discovery'; +import { + Discoveries, + RestDiscoveries, + DiscoveryPollingFactory as DiscoveryPollingFactory, + DefaultDiscoveryPollingFactory as DefaultDiscoveryPollingFactory +} from './Discovery'; import { container, Lifecycle } from 'tsyringe'; container @@ -167,6 +172,13 @@ container }, { lifecycle: Lifecycle.Singleton } ) + .register( + DiscoveryPollingFactory, + { + useClass: DefaultDiscoveryPollingFactory + }, + { lifecycle: Lifecycle.Singleton } + ) .register(Scans, { useClass: RestScans }, { lifecycle: Lifecycle.Singleton }) .register( Discoveries, diff --git a/src/index.ts b/src/index.ts index 63b2da73..27f3a41f 100644 --- a/src/index.ts +++ b/src/index.ts @@ -19,6 +19,7 @@ import container from './container'; import { RunDiscovery } from './Commands/RunDiscovery'; import { StopDiscovery } from './Commands/StopDiscovery'; import { RerunDiscovery } from './Commands/RerunDiscovery'; +import { PollingDiscoveryStatus } from './Commands/PollingDiscoveryStatus'; container.resolve(CliBuilder).build({ commands: [ @@ -31,6 +32,7 @@ container.resolve(CliBuilder).build({ new RunDiscovery(), new StopDiscovery(), new RerunDiscovery(), + new PollingDiscoveryStatus(), new UploadArchive(), new Configure(), new GetEntryPoints()