From d3e99435052aa676464d574dd75c8934e29abaf2 Mon Sep 17 00:00:00 2001 From: ztimson Date: Tue, 22 Mar 2022 22:31:30 -0400 Subject: [PATCH] Arg parser 2 --- scripts/botnet-manager.js | 133 +++++++++++++++++++------------------ scripts/lib/arg-parser2.js | 108 ++++++++++++++++++++++++++++++ 2 files changed, 175 insertions(+), 66 deletions(-) create mode 100644 scripts/lib/arg-parser2.js diff --git a/scripts/botnet-manager.js b/scripts/botnet-manager.js index 5fb1dd2..396cef5 100644 --- a/scripts/botnet-manager.js +++ b/scripts/botnet-manager.js @@ -1,4 +1,4 @@ -import {ArgError, ArgParser} from '/scripts/lib/arg-parser'; +import {ArgError, ArgParser} from '/scripts/lib/arg-parser2'; import {Logger} from '/scripts/lib/logger'; import {copyWithDependencies} from '/scripts/lib/utils'; @@ -108,79 +108,80 @@ class Manager { export async function main(ns) { // Setup ns.disableLog('ALL'); + const hostname = ns.getHostname(), portNum = 1; const argParser = new ArgParser('botnet-manager.js', 'Connect & manage a network of devices to launch distributed attacks.', [ - 'COPY [--HELP] [OPTIONS] FILE [DEST]', - 'JOIN [--HELP] MANAGER [DEVICE]', - 'KILL [--HELP]', - 'LEAVE [--HELP]', - 'MINE [--HELP] [OPTIONS] DEVICE', - 'RUN [--HELP] [OPTIONS] SCRIPT [ARGS]...', - 'START [--HELP] [OPTIONS]' - ], [ - new ArgParser('copy', 'Copy file & dependencies to swarm nodes', null, [ - {name: 'file', desc: 'File to copy', type: 'bool'}, - {name: 'dest', desc: 'File destination on nodes', optional: true, type: 'bool'}, - {name: 'manager', desc: 'Copy to manager node', flags: ['-m', '--manager'], type: 'bool'}, - {name: 'noDeps', desc: 'Skip copying dependencies', flags: ['-d', '--no-deps'], type: 'bool'}, - {name: 'workers', desc: 'Copy to worker nodes', flags: ['-w', '--workers'], type: 'bool'}, + new ArgParser('copy', 'Copy file & dependencies to swarm nodes', [ + {name: 'file', desc: 'File to copy', default: false}, + {name: 'manager', desc: 'Copy to manager node', flags: ['-m', '--manager'], default: false}, + {name: 'noDeps', desc: 'Skip copying dependencies', flags: ['-d', '--no-deps'], default: false}, + {name: 'workers', desc: 'Copy to worker nodes', flags: ['-w', '--workers'], default: false}, ]), - new ArgParser('join', 'Connect device as a worker node to the swarm', null, [ - {name: 'device', desc: 'Device to connect, defaults to the current machine', optional: true, default: ns.getHostname(), type: 'string'} + new ArgParser('join', 'Connect device as a worker node to the swarm', [ + {name: 'device', desc: 'Device to connect, defaults to the current machine', optional: true, default: hostname} ]), new ArgParser('kill', 'Kill any scripts running on worker nodes'), - new ArgParser('leave', 'Disconnect worker node from swarm', null, [ - {name: 'device', desc: 'Device to disconnect, defaults to the current machine', optional: true, default: ns.getHostname(), type: 'string'} + new ArgParser('leave', 'Disconnect worker node from swarm', [ + {name: 'device', desc: 'Device to disconnect, defaults to the current machine', optional: true, default: hostname} ]), - new ArgParser('run', 'Copy & run script on all worker nodes', null, [ + new ArgParser('run', 'Copy & run script on all worker nodes', [ {name: 'script', desc: 'Script to copy & execute', type: 'string'}, - {name: 'args', desc: 'Arguments for script. Forward the current target with: {{TARGET}}', optional: true, extras: true, type: 'string'}, + {name: 'args', desc: 'Arguments for script. Forward the current target with: {{TARGET}}', optional: true, extras: true}, ]), - new ArgParser('start', 'Start this device as the swarm manager') + new ArgParser('start', 'Start this device as the swarm manager'), + {name: 'silent', desc: 'Suppress program output', flags: ['-s', '--silent'], default: false}, ]); + const args = argParser.parse(ns.args); - try { - // Run - const portNum = 1; - const args = argParser.parse(ns.args); - if(args['command'].toLowerCase() == 'start') { // Start swarm manager - ns.tprint(`Starting swarm manager: ${args['remote']}`); - ns.tprint(`Connect a worker with: run swarm.js --join ${args['remote']}`); - await new Manager(ns, ns.getHostname(), portNum).start(); - } else { // Send a command to the swarm - if(args['command'] == 'copy') { - await this.ns.writePort(portNum, JSON.stringify({ - manager: args['remote'], - command: 'copy', - value: args['file'] - })); - } else if(args['command'] == 'join') { - await this.ns.writePort(portNum, JSON.stringify({ - manager: args['remote'], - command: 'join', - value: args['device'] - })); - } else if(args['command'] == 'kill') { - await this.ns.writePort(portNum, JSON.stringify({ - manager: args['remote'], - command: 'kill' - })); - } else if(args['command'] == 'leave') { - await this.ns.writePort(portNum, JSON.stringify({ - manager: args['remote'], - command: 'leave', - value: args['device'] - })); - } else if(args['command'] == 'run') { - await this.ns.writePort(portNum, JSON.stringify({ - manager: args['remote'], - command: 'run', - value: args['script'], - args: args['args'] - })); - } - } - } catch(err) { - if(err instanceof ArgError) return ns.tprint(parser.help(err.message)); - throw err; + // Help + if(args['help'] || args['_error']) + ns.tprint(argParser.help(args['help'] ? null : args['_error'], args['_command'])); + + // Run + 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']}`); + 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'], + command: 'copy', + value: args['copy']['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'], + command: 'join', + value: args['join']['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'], + 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'], + command: 'leave', + value: args['leave']['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'], + command: 'run', + value: args['run']['script'], + args: args['run']['args'] + })); } } diff --git a/scripts/lib/arg-parser2.js b/scripts/lib/arg-parser2.js new file mode 100644 index 0000000..3aa0519 --- /dev/null +++ b/scripts/lib/arg-parser2.js @@ -0,0 +1,108 @@ +export class ArgError extends Error {} + +export class ArgParser { + /** + * Create a unix-like argument parser to extract flags from the argument list. Can also create help messages. + * @param name {string} - Script name + * @param desc {string} - Help description + * @param argList {(ArgParser || {name: string, desc: string, flags: string[], optional: boolean, default: boolean})[]} - Array of CLI arguments + * @param examples {string[]} - Additional examples to display + */ + constructor(name, desc, argList = [], examples = []) { + this.name = name; + this.desc = desc; + + // Arguments + this.commands = argList.filter(arg => arg instanceof ArgParser); + this.args = argList.filter(arg => !arg.flags || !arg.flags.length); + this.flags = argList.filter(arg => !(arg instanceof ArgParser) && arg.flags && arg.flags.length); + this.flags.push({name: 'help', desc: 'Display this help message', flags: ['-h', '--help'], default: false}); + this.defaults = argList.reduce((acc, arg) => ({...acc, [arg.name]: arg?.extras ? [] : arg.default ?? null}), {}); + + // Examples + this.examples = [ + ...examples, + `[OPTIONS] ${this.args.map(arg => (arg.optional ? `[${arg.name.toUpperCase()}]` : arg.name.toUpperCase()) + (arg.extras ? '...' : '')).join(' ')}`, + this.commands.length ? `[OPTIONS] COMMAND` : null, + `--help ${this.commands.length ? '[COMMAND]' : ''}` + ].filter(e => !!e); + } + + /** + * Parse an array into an arguments dictionary using the configuration. + * @param args {string[]} - Array of arguments to be parsed + * @returns {object} - Dictionary of arguments with defaults applied + */ + parse(args) { + // Parse arguments + let extras = [], parsed = {...this.defaults}, queue = [...args]; + while(queue.length) { + let arg = queue.splice(0, 1)[0]; + if(arg[0] == '-') { // Flags + // Check for combined shorthand + if(arg[1] != '-' && arg.length > 2) { + queue = [...arg.substring(2).split('').map(a => `-${a}`), ...queue]; + arg = `-${arg[1]}`; + } + // 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` + 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)); + } 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`; + parsed[arg.name] = extras.splice(0, 1)[0]; + }); + // Extras + const extraKey = this.args.find(arg => arg.extras)?.name || '_extra'; + parsed[extraKey] = extras; + return parsed; + } + + /** + * Create help message from the provided description & argument list. + * @param message {string} - Message to display, defaults to the description + * @param command {string} - Command help message to show + * @returns {string} - Help message + */ + help(message = '', command = '') { + const spacer = (text) => Array(24 - text.length || 1).fill(' ').join(''); + + // Help with specific command + if(command) { + const argParser = this.commands.find(parser => parser.name == command); + if(!argParser) throw new Error(`${command.toUpperCase()} does not have a help`) + return argParser.help(message); + } + + // Description + let msg = `\n\n${message || this.desc}`; + // Examples + msg += '\n\nUsage:\t' + this.examples.map(ex => `run ${this.name} ${ex}`).join('\n\t'); + // Arguments + if(this.args.length) msg += '\n\n\t' + this.args + .map(arg => `${arg.name.toUpperCase()}${spacer(arg.name)}${arg.desc}`) + .join('\n\t'); + // Flags + msg += '\n\nOptions:\n\t' + this.flags.map(flag => { + const flags = flag.flags.join(', '); + return `${flags}${spacer(flags)}${flag.desc}`; + }).join('\n\t'); + // Commands + if(this.commands.length) msg += '\n\nCommands:\n\t' + this.commands + .map(command => `${command.name}${spacer(command.name)}${command.desc}`) + .join('\n\t'); + return `${msg}\n\n`; + } +}