From 98fe91ad36b846778756d5e4a74791b42014b2a6 Mon Sep 17 00:00:00 2001 From: ztimson Date: Thu, 24 Mar 2022 00:34:09 -0400 Subject: [PATCH] ArgParer2 --- scripts/botnet-manager.js | 164 ++++++++++++++++++------------------- scripts/lib/arg-parser2.js | 23 ++++-- scripts/lib/logger.js | 20 ++++- 3 files changed, 115 insertions(+), 92 deletions(-) diff --git a/scripts/botnet-manager.js b/scripts/botnet-manager.js index 396cef5..4ac9e4b 100644 --- a/scripts/botnet-manager.js +++ b/scripts/botnet-manager.js @@ -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 {copyWithDependencies} from '/scripts/lib/utils'; @@ -14,53 +14,67 @@ class Manager { () => `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` ]); - 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() { const state = JSON.parse(await this.ns.read(this.config) || 'null'); if(state) { this.running = state.running; - if(this.running) await this.runCommand(this.running['command'], this.running); this.workers = state.workers; + if(this.running) await this.runCommand(this.running); } } - async runCommand(command, extra = null) { - if (command == 'copy') { - this.logger.log(`Copying: ${extra['value']}`); - await this.workerExec(async w => await copyWithDependencies(ns, extra['value'], w)); - } else if (command == 'join') { - const exists = this.workers.findIndex(w => w.hostname == extra['value']); - this.workers.splice(exists, 1); - this.workers.push(this.ns.getServer(extra['value'])); - this.logger.log(`${exists != -1 ? 'Reconnected' : 'Connected:'}: ${extra['value']}`); - if(this.running) await this.runCommand(this.running['command'], {...this.running, device: extra['value']}); - } else if (command == 'kill') { - this.logger.log('Killing scripts'); - await this.workerExec(w => this.ns.killall(w.hostname)); - this.running = null; - } else if (command == 'leave') { - this.logger.log(`Disconnecting: ${extra['value']}`); - const worker = this.workers.splice(this.workers.findIndex(w => w.hostname == extra['value']), 1); - this.ns.killall(worker.hostname); - } else if (command == 'run') { - await this.runCommand('copy', {value: extra['value']}); - await this.runCommand('kill'); - const run = (w) => { - const threads = ~~(w.maxRam / this.ns.getScriptRam(extra['value'], this.ns.getHostname())) || 1; - this.ns.exec(extra['value'], w, threads, ...(extra['args'] || [])); - } - if(extra['device']) { - const w = this.workers.find(w => w.hostname == extra['device']); - if(w) run(w); - } else { - this.logger.log(`Starting script: ${extra['value']}`); - await this.workerExec(run); - this.running = {[command]: command, ...extra}; + async runCommand(request = null) { + try { + if (request['command'] == 'copy') { + this.logger.log(`Copying: ${request['value']}`); + await this.workerExec(async w => await copyWithDependencies(ns, request['value'], w)); + } else if (request['command'] == 'join') { + if(request['value'] == this.device) throw 'Cannot connect manager as a worker'; + const exists = this.workers.findIndex(w => w.hostname == request['value']); + if(exists != -1) this.workers.splice(exists, 1); + this.workers.push(this.ns.getServer(request['value'])); + this.logger.log(`${exists != -1 ? 'Reconnected' : 'Connected'}: ${request['value']}`); + if(this.running) await this.runCommand({ + ...this.running, + device: request['value'] + }); + } else if (request['command'] == 'kill') { + this.logger.log('Killing scripts'); + await this.workerExec(w => this.ns.killall(w.hostname)); + this.running = null; + } else if (request['command'] == 'leave') { + this.logger.log(`Disconnecting: ${request['value']}`); + const worker = this.workers.splice(this.workers.findIndex(w => w.hostname == request['value']), 1); + this.ns.killall(worker.hostname); + } else if (request['command'] == 'run') { + await this.runCommand('copy', {value: request['value']}); + await this.runCommand('kill'); + const run = (w) => { + const threads = ~~(w.maxRam / this.ns.getScriptRam(request['value'], this.ns.getHostname())) || 1; + this.ns.exec(request['value'], w, threads, ...(request['args'] || [])); + } + if (request['device']) { + 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'); } - async start(load = true) { - if(load) await this.load(); - let checkTick = -1, runCheck = false; - for(let tick = 1; true; tick = tick == 3600 ? 1 : tick + 1) { - if(tick == checkTick) runCheck = true; + async start() { + if(this.config) await this.load(); + let checkTick = 3600, runCheck = false; + for(let tick = 1; true; tick = tick > 3600 ? 1 : tick + 1) { await this.ns.sleep(1000); + this.logger.log(); + if(tick == checkTick) runCheck = true; + let req = this.rx(); - // Check for new commands - const payload = this.ns.readPort(this.port); - - // Check if we need to update the running command every hour - if(payload == 'NULL PORT DATA' && runCheck && this.running['update']) { - runCheck = false; - await this.runCommand(this.running['command'], this.running); + // Check if we are idle + if(!req) { + // Check if we need to update the running command while we are idle + if(runCheck && this.running['update']) { + runCheck = false; + await this.runCommand(this.running['command'], this.running); + } continue; } // Run command - if(this.isCommand(payload)) { - checkTick = tick; - await this.runCommand(payload['command'], payload); + if((req = JSON.parse(req)) && this.isCommand(req)) { + await this.runCommand(req); await this.save(); } 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 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['start']['help'] || args['start']['_error']) - ns.tprint(argParser.help(args['start']['help'] ? null : args['start']['_error'], 'start')); - ns.tprint(`Starting swarm manager: ${args['remote']}`); - ns.tprint(`Connect a worker with: run swarm.js --join ${args['remote']}`); + ns.tprint(`Starting swarm manager: ${hostname}`); + ns.tprint(`Connect a worker with: run botnet-manager.js --join ${hostname}`); await new Manager(ns, hostname, portNum).start(); } else if(args['_command'] == 'copy') { // Issue copy command - if(args['copy']['help'] || args['copy']['_error']) - ns.tprint(argParser.help(args['copy']['help'] ? null : args['copy']['_error'], 'copy')); - await this.ns.writePort(portNum, JSON.stringify({ - manager: args['copy']['remote'], + await ns.writePort(portNum, JSON.stringify({ command: 'copy', - value: args['copy']['file'] + value: args['file'] })); } else if(args['_command'] == 'join') { // Issue join command - if(args['join']['help'] || args['join']['_error']) - ns.tprint(argParser.help(args['join']['help'] ? null : args['join']['_error'], 'join')); - await this.ns.writePort(portNum, JSON.stringify({ - manager: args['join']['remote'], + await ns.writePort(portNum, JSON.stringify({ command: 'join', - value: args['join']['device'] + value: args['device'] })); } else if(args['_command'] == 'kill') { // Issue kill command - if(args['kill']['help'] || args['kill']['_error']) - ns.tprint(argParser.help(args['kill']['help'] ? null : args['kill']['_error'], 'kill')); - await this.ns.writePort(portNum, JSON.stringify({ - manager: args['kill']['remote'], + await ns.writePort(portNum, JSON.stringify({ command: 'kill' })); } else if(args['_command'] == 'leave') { // Issue leave command - if(args['leave']['help'] || args['leave']['_error']) - ns.tprint(argParser.help(args['leave']['help'] ? null : args['leave']['_error'], 'leave')); - await this.ns.writePort(portNum, JSON.stringify({ - manager: args['leave']['remote'], + await ns.writePort(portNum, JSON.stringify({ command: 'leave', - value: args['leave']['device'] + value: args['device'] })); } else if(args['_command'] == 'run') { // Issue run command - if(args['run']['help'] || args['run']['_error']) - ns.tprint(argParser.help(args['run']['help'] ? null : args['run']['_error'], 'run')); - await this.ns.writePort(portNum, JSON.stringify({ - manager: args['run']['remote'], + await ns.writePort(portNum, JSON.stringify({ command: 'run', - value: args['run']['script'], - args: args['run']['args'] + value: args['script'], + args: args['args'] })); } } diff --git a/scripts/lib/arg-parser2.js b/scripts/lib/arg-parser2.js index 3aa0519..854800d 100644 --- a/scripts/lib/arg-parser2.js +++ b/scripts/lib/arg-parser2.js @@ -47,21 +47,32 @@ export class ArgParser { // Find & add flag const combined = arg.split('='); const argDef = this.flags.find(flag => flag.flags.includes(combined[0] || arg)); - if(!argDef) extras.push(arg); // Not found, add to extras - const value = argDef.default === false ? true : argDef.default === true ? false : queue.splice(queue.findIndex(q => q[0] != '-'), 1)[0]; - if(value == null) parsed['_error'] = `${argDef.name.toUpperCase()} missing value` + if(argDef == null) { // Not found, add to extras + extras.push(arg); + 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; } else { // Command const c = this.commands.find(command => command.name == arg); if(!!c) { - parsed['_command'] = c.name; - parsed[c.name] = c.parse(queue.splice(0, queue.length)); + const parsedCommand = 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 } } // Arguments 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]; }); // Extras diff --git a/scripts/lib/logger.js b/scripts/lib/logger.js index 8049831..c95094b 100644 --- a/scripts/lib/logger.js +++ b/scripts/lib/logger.js @@ -12,8 +12,15 @@ export class Logger { this.fns = lineFns; this.historyLen -= lineFns.length * 2; 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 */ @@ -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.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.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}`); } }