Arg parser 2

This commit is contained in:
Zakary Timson 2022-03-22 22:31:30 -04:00
parent 44af8bb9a2
commit d3e9943505
2 changed files with 175 additions and 66 deletions

View File

@ -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},
]);
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();
} else { // Send a command to the swarm
if(args['command'] == 'copy') {
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['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'],
args: args['args']
value: args['run']['script'],
args: args['run']['args']
}));
}
}
} catch(err) {
if(err instanceof ArgError) return ns.tprint(parser.help(err.message));
throw err;
}
}

108
scripts/lib/arg-parser2.js Normal file
View File

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