Arg parser 2

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]',
+        new ArgParser('copy', 'Copy file & dependencies to swarm nodes', [
-        'JOIN [--HELP] MANAGER [DEVICE]',
+            {name: 'file', desc: 'File to copy', default: false},
-        'KILL [--HELP]',
+            {name: 'manager', desc: 'Copy to manager node', flags: ['-m', '--manager'], default: false},
-        'LEAVE [--HELP]',
+            {name: 'noDeps', desc: 'Skip copying dependencies', flags: ['-d', '--no-deps'], default: false},
-        'MINE [--HELP] [OPTIONS] DEVICE',
+            {name: 'workers', desc: 'Copy to worker nodes', flags: ['-w', '--workers'], default: false},
-        '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('join', 'Connect device as a worker node to the swarm', null, [
+        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: ns.getHostname(), type: 'string'}
+            {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, [
+        new ArgParser('leave', 'Disconnect worker node from swarm', [
-            {name: 'device', desc: 'Device to disconnect, defaults to the current machine', optional: true, default: ns.getHostname(), type: 'string'}
+            {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},
     ]);
-
-    try {
-        // Run
-        const portNum = 1;
     const args = argParser.parse(ns.args);
-        if(args['command'].toLowerCase() == 'start') { // Start swarm manager
+
+    // 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, ns.getHostname(), portNum).start();
+        await new Manager(ns, hostname, portNum).start();
-        } else { // Send a command to the swarm
+    } else if(args['_command'] == 'copy') { // Issue copy command
-            if(args['command'] == 'copy') {
+        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['remote'],
+            manager: args['copy']['remote'],
             command: 'copy',
-                    value: args['file']
+            value: args['copy']['file']
         }));
-            } else if(args['command'] == 'join') {
+    } 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['remote'],
+            manager: args['join']['remote'],
             command: 'join',
-                    value: args['device']
+            value: args['join']['device']
         }));
-            } else if(args['command'] == 'kill') {
+    } 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['remote'],
+            manager: args['kill']['remote'],
             command: 'kill'
         }));
-            } else if(args['command'] == 'leave') {
+    } 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['remote'],
+            manager: args['leave']['remote'],
             command: 'leave',
-                    value: args['device']
+            value: args['leave']['device']
         }));
-            } else if(args['command'] == 'run') {
+    } 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['remote'],
+            manager: args['run']['remote'],
             command: 'run',
-                    value: args['script'],
+            value: args['run']['script'],
-                    args: args['args']
+            args: args['run']['args']
         }));
     }
-        }
-    } catch(err) {
-        if(err instanceof ArgError) return ns.tprint(parser.help(err.message));
-        throw err;
-    }
 }

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`;
+    }
+}