ArgParer2
This commit is contained in:
		| @@ -1,4 +1,4 @@ | |||||||
| import {ArgError, ArgParser} from '/scripts/lib/arg-parser2'; | import {ArgParser} from '/scripts/lib/arg-parser2'; | ||||||
| import {Logger} from '/scripts/lib/logger'; | import {Logger} from '/scripts/lib/logger'; | ||||||
| import {copyWithDependencies} from '/scripts/lib/utils'; | import {copyWithDependencies} from '/scripts/lib/utils'; | ||||||
|  |  | ||||||
| @@ -14,53 +14,67 @@ class Manager { | |||||||
|             () => `Swarm Manager: ${device}`, |             () => `Swarm Manager: ${device}`, | ||||||
|             () => `Workers: ${this.workers.length}\tCores: ${this.workers.reduce((acc, w) => acc + w.cpuCores, 0)}\tRAM: ${this.workers.reduce((acc, w) => acc + w.maxRam, 0)} GB` |             () => `Workers: ${this.workers.length}\tCores: ${this.workers.reduce((acc, w) => acc + w.cpuCores, 0)}\tRAM: ${this.workers.reduce((acc, w) => acc + w.maxRam, 0)} GB` | ||||||
|         ]); |         ]); | ||||||
|         this.port = port; |  | ||||||
|  |         // Port communication | ||||||
|  |         this.rx = () => { | ||||||
|  |             let payload = ns.readPort(port); | ||||||
|  |             return payload == 'NULL PORT DATA' ? null : payload; | ||||||
|  |         } | ||||||
|  |         this.tx = (payload) => ns.writePort(portNum + 10, JSON.stringify(payload)); | ||||||
|     } |     } | ||||||
|  |  | ||||||
|     isCommand(payload) { return payload['manager'] == this.device && payload['command'] != null; } |     isCommand(payload) { return ['copy', 'join', 'kill', 'leave', 'run'].includes(payload['command']); } | ||||||
|  |  | ||||||
|     async load() { |     async load() { | ||||||
|         const state = JSON.parse(await this.ns.read(this.config) || 'null'); |         const state = JSON.parse(await this.ns.read(this.config) || 'null'); | ||||||
|         if(state) { |         if(state) { | ||||||
|             this.running = state.running; |             this.running = state.running; | ||||||
|             if(this.running) await this.runCommand(this.running['command'], this.running); |  | ||||||
|             this.workers = state.workers; |             this.workers = state.workers; | ||||||
|  |             if(this.running) await this.runCommand(this.running); | ||||||
|         } |         } | ||||||
|     } |     } | ||||||
|  |  | ||||||
|     async runCommand(command, extra = null) { |     async runCommand(request = null) { | ||||||
|         if (command == 'copy') { |         try { | ||||||
|             this.logger.log(`Copying: ${extra['value']}`); |             if (request['command'] == 'copy') { | ||||||
|             await this.workerExec(async w => await copyWithDependencies(ns, extra['value'], w)); |                 this.logger.log(`Copying: ${request['value']}`); | ||||||
|         } else if (command == 'join') { |                 await this.workerExec(async w => await copyWithDependencies(ns, request['value'], w)); | ||||||
|             const exists = this.workers.findIndex(w => w.hostname == extra['value']); |             } else if (request['command'] == 'join') { | ||||||
|             this.workers.splice(exists, 1); |                 if(request['value'] == this.device) throw 'Cannot connect manager as a worker'; | ||||||
|             this.workers.push(this.ns.getServer(extra['value'])); |                 const exists = this.workers.findIndex(w => w.hostname == request['value']); | ||||||
|             this.logger.log(`${exists != -1 ? 'Reconnected' : 'Connected:'}: ${extra['value']}`); |                 if(exists != -1) this.workers.splice(exists, 1); | ||||||
|             if(this.running) await this.runCommand(this.running['command'], {...this.running, device: extra['value']}); |                 this.workers.push(this.ns.getServer(request['value'])); | ||||||
|         } else if (command == 'kill') { |                 this.logger.log(`${exists != -1 ? 'Reconnected' : 'Connected'}: ${request['value']}`); | ||||||
|             this.logger.log('Killing scripts'); |                 if(this.running) await this.runCommand({ | ||||||
|             await this.workerExec(w => this.ns.killall(w.hostname)); |                     ...this.running, | ||||||
|             this.running = null; |                     device: request['value'] | ||||||
|         } else if (command == 'leave') { |                 }); | ||||||
|             this.logger.log(`Disconnecting: ${extra['value']}`); |             } else if (request['command'] == 'kill') { | ||||||
|             const worker = this.workers.splice(this.workers.findIndex(w => w.hostname == extra['value']), 1); |                 this.logger.log('Killing scripts'); | ||||||
|             this.ns.killall(worker.hostname); |                 await this.workerExec(w => this.ns.killall(w.hostname)); | ||||||
|         } else if (command == 'run') { |                 this.running = null; | ||||||
|             await this.runCommand('copy', {value: extra['value']}); |             } else if (request['command'] == 'leave') { | ||||||
|             await this.runCommand('kill'); |                 this.logger.log(`Disconnecting: ${request['value']}`); | ||||||
|             const run = (w) => { |                 const worker = this.workers.splice(this.workers.findIndex(w => w.hostname == request['value']), 1); | ||||||
|                 const threads = ~~(w.maxRam / this.ns.getScriptRam(extra['value'], this.ns.getHostname())) || 1; |                 this.ns.killall(worker.hostname); | ||||||
|                 this.ns.exec(extra['value'], w, threads, ...(extra['args'] || [])); |             } else if (request['command'] == 'run') { | ||||||
|             } |                 await this.runCommand('copy', {value: request['value']}); | ||||||
|             if(extra['device']) { |                 await this.runCommand('kill'); | ||||||
|                 const w = this.workers.find(w => w.hostname == extra['device']); |                 const run = (w) => { | ||||||
|                 if(w) run(w); |                     const threads = ~~(w.maxRam / this.ns.getScriptRam(request['value'], this.ns.getHostname())) || 1; | ||||||
|             } else { |                     this.ns.exec(request['value'], w, threads, ...(request['args'] || [])); | ||||||
|                 this.logger.log(`Starting script: ${extra['value']}`); |                 } | ||||||
|                 await this.workerExec(run); |                 if (request['device']) { | ||||||
|                 this.running = {[command]: command, ...extra}; |                     const w = this.workers.find(w => w.hostname == request['device']); | ||||||
|  |                     if (w) run(w); | ||||||
|  |                 } else { | ||||||
|  |                     this.logger.log(`Starting script: ${request['value']}`); | ||||||
|  |                     await this.workerExec(run); | ||||||
|  |                     this.running = request; | ||||||
|  |                 } | ||||||
|             } |             } | ||||||
|  |         } catch (e) { | ||||||
|  |             this.logger.error(`${request['command']} - ${e}`); | ||||||
|         } |         } | ||||||
|     } |     } | ||||||
|  |  | ||||||
| @@ -71,30 +85,31 @@ class Manager { | |||||||
|         }), 'w'); |         }), 'w'); | ||||||
|     } |     } | ||||||
|  |  | ||||||
|     async start(load = true) { |     async start() { | ||||||
|         if(load) await this.load(); |         if(this.config) await this.load(); | ||||||
|         let checkTick = -1, runCheck = false; |         let checkTick = 3600, runCheck = false; | ||||||
|         for(let tick = 1; true; tick = tick == 3600 ? 1 : tick + 1) { |         for(let tick = 1; true; tick = tick > 3600 ? 1 : tick + 1) { | ||||||
|             if(tick == checkTick) runCheck = true; |  | ||||||
|             await this.ns.sleep(1000); |             await this.ns.sleep(1000); | ||||||
|  |             this.logger.log(); | ||||||
|  |             if(tick == checkTick) runCheck = true; | ||||||
|  |             let req = this.rx(); | ||||||
|  |  | ||||||
|             // Check for new commands |             // Check if we are idle | ||||||
|             const payload = this.ns.readPort(this.port); |             if(!req) { | ||||||
|  |                 // Check if we need to update the running command while we are idle | ||||||
|             // Check if we need to update the running command every hour |                 if(runCheck && this.running['update']) { | ||||||
|             if(payload == 'NULL PORT DATA' && runCheck && this.running['update']) { |                     runCheck = false; | ||||||
|                 runCheck = false; |                     await this.runCommand(this.running['command'], this.running); | ||||||
|                 await this.runCommand(this.running['command'], this.running); |                 } | ||||||
|                 continue; |                 continue; | ||||||
|             } |             } | ||||||
|  |  | ||||||
|             // Run command |             // Run command | ||||||
|             if(this.isCommand(payload)) { |             if((req = JSON.parse(req)) && this.isCommand(req)) { | ||||||
|                 checkTick = tick; |                 await this.runCommand(req); | ||||||
|                 await this.runCommand(payload['command'], payload); |  | ||||||
|                 await this.save(); |                 await this.save(); | ||||||
|             } else { // Invalid command |             } else { // Invalid command | ||||||
|                 this.logger.log(`Unknown command: ${JSON.stringify(payload)}`); |                 this.logger.log(`Unknown command: ${JSON.stringify(req)}`); | ||||||
|             } |             } | ||||||
|         } |         } | ||||||
|     } |     } | ||||||
| @@ -134,54 +149,37 @@ export async function main(ns) { | |||||||
|  |  | ||||||
|     // Help |     // Help | ||||||
|     if(args['help'] || args['_error']) |     if(args['help'] || args['_error']) | ||||||
|         ns.tprint(argParser.help(args['help'] ? null : args['_error'], args['_command'])); |         return ns.tprint(argParser.help(args['help'] ? null : args['_error'], args['_command'])); | ||||||
|  |  | ||||||
|     // Run |     // Run command | ||||||
|     if(args['_command'] == 'start') { // Start botnet manager |     if(args['_command'] == 'start') { // Start botnet manager | ||||||
|         if(args['start']['help'] || args['start']['_error']) |         ns.tprint(`Starting swarm manager: ${hostname}`); | ||||||
|             ns.tprint(argParser.help(args['start']['help'] ? null : args['start']['_error'], 'start')); |         ns.tprint(`Connect a worker with: run botnet-manager.js --join ${hostname}`); | ||||||
|         ns.tprint(`Starting swarm manager: ${args['remote']}`); |  | ||||||
|         ns.tprint(`Connect a worker with: run swarm.js --join ${args['remote']}`); |  | ||||||
|         await new Manager(ns, hostname, portNum).start(); |         await new Manager(ns, hostname, portNum).start(); | ||||||
|     } else if(args['_command'] == 'copy') { // Issue copy command |     } else if(args['_command'] == 'copy') { // Issue copy command | ||||||
|         if(args['copy']['help'] || args['copy']['_error']) |         await ns.writePort(portNum, JSON.stringify({ | ||||||
|             ns.tprint(argParser.help(args['copy']['help'] ? null : args['copy']['_error'], 'copy')); |  | ||||||
|         await this.ns.writePort(portNum, JSON.stringify({ |  | ||||||
|             manager: args['copy']['remote'], |  | ||||||
|             command: 'copy', |             command: 'copy', | ||||||
|             value: args['copy']['file'] |             value: args['file'] | ||||||
|         })); |         })); | ||||||
|     } else if(args['_command'] == 'join') { // Issue join command |     } else if(args['_command'] == 'join') { // Issue join command | ||||||
|         if(args['join']['help'] || args['join']['_error']) |         await ns.writePort(portNum, JSON.stringify({ | ||||||
|             ns.tprint(argParser.help(args['join']['help'] ? null : args['join']['_error'], 'join')); |  | ||||||
|         await this.ns.writePort(portNum, JSON.stringify({ |  | ||||||
|             manager: args['join']['remote'], |  | ||||||
|             command: 'join', |             command: 'join', | ||||||
|             value: args['join']['device'] |             value: args['device'] | ||||||
|         })); |         })); | ||||||
|     } else if(args['_command'] == 'kill') { // Issue kill command |     } else if(args['_command'] == 'kill') { // Issue kill command | ||||||
|         if(args['kill']['help'] || args['kill']['_error']) |         await ns.writePort(portNum, JSON.stringify({ | ||||||
|             ns.tprint(argParser.help(args['kill']['help'] ? null : args['kill']['_error'], 'kill')); |  | ||||||
|         await this.ns.writePort(portNum, JSON.stringify({ |  | ||||||
|             manager: args['kill']['remote'], |  | ||||||
|             command: 'kill' |             command: 'kill' | ||||||
|         })); |         })); | ||||||
|     } else if(args['_command'] == 'leave') { // Issue leave command |     } else if(args['_command'] == 'leave') { // Issue leave command | ||||||
|         if(args['leave']['help'] || args['leave']['_error']) |         await ns.writePort(portNum, JSON.stringify({ | ||||||
|             ns.tprint(argParser.help(args['leave']['help'] ? null : args['leave']['_error'], 'leave')); |  | ||||||
|         await this.ns.writePort(portNum, JSON.stringify({ |  | ||||||
|             manager: args['leave']['remote'], |  | ||||||
|             command: 'leave', |             command: 'leave', | ||||||
|             value: args['leave']['device'] |             value: args['device'] | ||||||
|         })); |         })); | ||||||
|     } else if(args['_command'] == 'run') { // Issue run command |     } else if(args['_command'] == 'run') { // Issue run command | ||||||
|         if(args['run']['help'] || args['run']['_error']) |         await ns.writePort(portNum, JSON.stringify({ | ||||||
|             ns.tprint(argParser.help(args['run']['help'] ? null : args['run']['_error'], 'run')); |  | ||||||
|         await this.ns.writePort(portNum, JSON.stringify({ |  | ||||||
|             manager: args['run']['remote'], |  | ||||||
|             command: 'run', |             command: 'run', | ||||||
|             value: args['run']['script'], |             value: args['script'], | ||||||
|             args: args['run']['args'] |             args: args['args'] | ||||||
|         })); |         })); | ||||||
|     } |     } | ||||||
| } | } | ||||||
|   | |||||||
| @@ -47,21 +47,32 @@ export class ArgParser { | |||||||
|                 // Find & add flag |                 // Find & add flag | ||||||
|                 const combined = arg.split('='); |                 const combined = arg.split('='); | ||||||
|                 const argDef = this.flags.find(flag => flag.flags.includes(combined[0] || arg)); |                 const argDef = this.flags.find(flag => flag.flags.includes(combined[0] || arg)); | ||||||
|                 if(!argDef) extras.push(arg); // Not found, add to extras |                 if(argDef == null) { // Not found, add to extras | ||||||
|                 const value = argDef.default === false ? true : argDef.default === true ? false : queue.splice(queue.findIndex(q => q[0] != '-'), 1)[0]; |                     extras.push(arg); | ||||||
|                 if(value == null) parsed['_error'] = `${argDef.name.toUpperCase()} missing value` |                     continue; | ||||||
|  |                 } | ||||||
|  |                 const value = argDef.default === false ? true : argDef.default === true ? false : argDef.default || queue.splice(queue.findIndex(q => q[0] != '-'), 1)[0]; | ||||||
|  |                 if(value == null) parsed['_error'] = `Option missing value: ${arg.name}`; | ||||||
|                 parsed[argDef.name] = value; |                 parsed[argDef.name] = value; | ||||||
|             } else { // Command |             } else { // Command | ||||||
|                 const c = this.commands.find(command => command.name == arg); |                 const c = this.commands.find(command => command.name == arg); | ||||||
|                 if(!!c) { |                 if(!!c) { | ||||||
|                     parsed['_command'] = c.name; |                     const parsedCommand = c.parse(queue.splice(0, queue.length)); | ||||||
|                     parsed[c.name] = c.parse(queue.splice(0, queue.length)); |                     Object.keys(parsedCommand).forEach(key => { | ||||||
|  |                         if(parsed[key] != parsedCommand[key] && parsedCommand[key] == c.defaults[key]) | ||||||
|  |                             delete parsedCommand[key]; | ||||||
|  |                     }); | ||||||
|  |                     parsed = { | ||||||
|  |                         ...parsed, | ||||||
|  |                         ...parsedCommand, | ||||||
|  |                         _command: c.name | ||||||
|  |                     }; | ||||||
|                 } else extras.push(arg); // Not found, add to extras |                 } else extras.push(arg); // Not found, add to extras | ||||||
|             } |             } | ||||||
|         } |         } | ||||||
|         // Arguments |         // Arguments | ||||||
|         this.args.filter(arg => !arg.extras).forEach(arg => { |         this.args.filter(arg => !arg.extras).forEach(arg => { | ||||||
|             if(!arg.optional && !extras.length) parsed['_error'] = `${arg.name.toUpperCase()} is missing`; |             if(!arg.optional && !extras.length) parsed['_error'] = `Argument missing: ${arg.name.toUpperCase()}`; | ||||||
|             parsed[arg.name] = extras.splice(0, 1)[0]; |             parsed[arg.name] = extras.splice(0, 1)[0]; | ||||||
|         }); |         }); | ||||||
|         // Extras |         // Extras | ||||||
|   | |||||||
| @@ -12,8 +12,15 @@ export class Logger { | |||||||
| 		this.fns = lineFns; | 		this.fns = lineFns; | ||||||
| 		this.historyLen -= lineFns.length * 2; | 		this.historyLen -= lineFns.length * 2; | ||||||
| 		this.history = Array(this.historyLen).fill(''); | 		this.history = Array(this.historyLen).fill(''); | ||||||
|  | 		this.log(); | ||||||
| 	} | 	} | ||||||
|  |  | ||||||
|  | 	/** | ||||||
|  | 	 * Add red error message to logs | ||||||
|  | 	 * @param message {string} - Text that will be added | ||||||
|  | 	 */ | ||||||
|  | 	error(message) { this.log(`ERROR: ${message}`); } | ||||||
|  |  | ||||||
| 	/** | 	/** | ||||||
| 	 * Add a linebreak | 	 * Add a linebreak | ||||||
| 	 */ | 	 */ | ||||||
| @@ -33,13 +40,20 @@ export class Logger { | |||||||
| 	} | 	} | ||||||
|  |  | ||||||
| 	/** | 	/** | ||||||
| 	 * Add message to logs & output | 	 * Add message to the logs | ||||||
|  | 	 * @param message {string} - Text that will be added | ||||||
| 	 */ | 	 */ | ||||||
| 	log(message) { | 	log(message = '') { | ||||||
| 		this.ns.clearLog(); | 		this.ns.clearLog(); | ||||||
| 		this.header(); | 		this.header(); | ||||||
| 		if(message != null) this.history.push(message); | 		if(message) this.history.push(message); | ||||||
| 		this.history.splice(0, this.history.length - this.historyLen); | 		this.history.splice(0, this.history.length - this.historyLen); | ||||||
| 		this.history.forEach(m => this.ns.print(m)); | 		this.history.forEach(m => this.ns.print(m)); | ||||||
| 	} | 	} | ||||||
|  |  | ||||||
|  | 	/** | ||||||
|  | 	 * Add orange warning to the logs | ||||||
|  | 	 * @param message {string} - Text that will be added | ||||||
|  | 	 */ | ||||||
|  | 	warn(message) { this.log(`WARN: ${message}`); } | ||||||
| } | } | ||||||
|   | |||||||
		Reference in New Issue
	
	Block a user