ArgParer2

This commit is contained in:
Zakary Timson 2022-03-24 00:34:09 -04:00
parent d3e9943505
commit 98fe91ad36
3 changed files with 115 additions and 92 deletions

View File

@ -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,54 +14,68 @@ 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']}`);
if(this.running) await this.runCommand({
...this.running,
device: request['value']
});
} else if (request['command'] == 'kill') {
this.logger.log('Killing scripts'); this.logger.log('Killing scripts');
await this.workerExec(w => this.ns.killall(w.hostname)); await this.workerExec(w => this.ns.killall(w.hostname));
this.running = null; this.running = null;
} else if (command == 'leave') { } else if (request['command'] == 'leave') {
this.logger.log(`Disconnecting: ${extra['value']}`); this.logger.log(`Disconnecting: ${request['value']}`);
const worker = this.workers.splice(this.workers.findIndex(w => w.hostname == extra['value']), 1); const worker = this.workers.splice(this.workers.findIndex(w => w.hostname == request['value']), 1);
this.ns.killall(worker.hostname); this.ns.killall(worker.hostname);
} else if (command == 'run') { } else if (request['command'] == 'run') {
await this.runCommand('copy', {value: extra['value']}); await this.runCommand('copy', {value: request['value']});
await this.runCommand('kill'); await this.runCommand('kill');
const run = (w) => { const run = (w) => {
const threads = ~~(w.maxRam / this.ns.getScriptRam(extra['value'], this.ns.getHostname())) || 1; const threads = ~~(w.maxRam / this.ns.getScriptRam(request['value'], this.ns.getHostname())) || 1;
this.ns.exec(extra['value'], w, threads, ...(extra['args'] || [])); this.ns.exec(request['value'], w, threads, ...(request['args'] || []));
} }
if(extra['device']) { if (request['device']) {
const w = this.workers.find(w => w.hostname == extra['device']); const w = this.workers.find(w => w.hostname == request['device']);
if (w) run(w); if (w) run(w);
} else { } else {
this.logger.log(`Starting script: ${extra['value']}`); this.logger.log(`Starting script: ${request['value']}`);
await this.workerExec(run); await this.workerExec(run);
this.running = {[command]: command, ...extra}; this.running = request;
} }
} }
} catch (e) {
this.logger.error(`${request['command']} - ${e}`);
}
} }
async save() { async save() {
@ -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']
})); }));
} }
} }

View File

@ -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

View File

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