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

View File

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

View File

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