firmament-vita/ts/implementations/vita-tasks-impl.ts

Firmament module for automating VITA tasks

import {injectable, inject} from 'inversify';
import {CommandUtil, ForceErrorImpl, ProgressBar} from "firmament-yargs";
import {VitaFileUtil} from "../interfaces/vita-file-util";
import {VitaTasks} from "../interfaces/vita-tasks";
import {
  DecryptAndUnTarResult, UnZipFileResult, ProcessFullPipelineResult, MergePcapFilesResult, FullPipeline,
  ProcessFullPipelineStatus, DecryptAndUnTarStatus, UnZipFileStatus, MergePcapFilesStatus
} from "../interfaces/vita-options-results";
import {VitaDecryptUnTar} from "../interfaces/vita-decrypt-untar";
import {VitaPcapMerge} from "../interfaces/vita-pcap-merge";
import kernel from "../inversify.config";
import {EtlFlow, EtlError} from "etl-typings";

const _ = require('lodash');
const shorthash = require('shorthash');
const deepExtend = require('deep-extend');
const async = require('async');

@injectable()
export class VitaTasksImpl extends ForceErrorImpl implements VitaTasks {
  private count = 0;

  constructor(@inject('CommandUtil') private commandUtil: CommandUtil,
              @inject('VitaFileUtil') private vitaFileUtil: VitaFileUtil,
              @inject('ProgressBar') private progressBar: ProgressBar,
              @inject('VitaPcapMerge') private vitaPcapMerge: VitaPcapMerge,
              @inject('VitaDecryptUnTar') private vitaDecryptUnTar: VitaDecryptUnTar) {
    super();
  }

  jsonOut(argv: any) {
    let fullPipeline = kernel.get<FullPipeline>('FullPipeline');
    fullPipeline.fromArgv(argv);
    if (argv.write) {
      fullPipeline.toJsonFile(argv.write);
    } else {
      this.commandUtil.stdoutWrite(fullPipeline.toJson());
    }
    this.commandUtil.processExit();
  }

  private static hash(input: string): string {
    //const retVal = crypto.createHash('sha256').update(input).digest('base65').toString();
    let retVal: string = shorthash.unique(input);
    //Make sure it doesn't start with a number (has to be a valid JS identifier)
    if (!isNaN(parseInt(retVal[0]))) {
      retVal = `A${retVal}`;
    }
    return retVal;
  }

  process(argv: any) {
    let fullPipeline = kernel.get<FullPipeline>('FullPipeline');
    fullPipeline.fromArgv(argv);
    this.processFullPipelineInstance(fullPipeline,
      (err: Error, result: ProcessFullPipelineStatus) => {
        if (result.decryptAndUnTarStatus) {
          let {taskName, current, total} = result.decryptAndUnTarStatus;
          this.progressBar.showProgressForTask('decryptAndUnTarStatus', taskName, current, total);
        } else if (result.unZipFileStatus) {
          this.progressBar.showProgressForTask('decryptAndUnTarStatus', '', 1, 0);
          let {taskName, current, total} = result.unZipFileStatus;
          this.progressBar.showProgressForTask(VitaTasksImpl.hash(taskName), taskName, current, total);
        } else if (result.mergePcapFilesStatus) {
          let {taskName, current, total} = result.mergePcapFilesStatus;
          this.progressBar.showProgressForTask(VitaTasksImpl.hash(taskName), taskName, current, total);
        }
      },
      (err: Error, result: ProcessFullPipelineResult) => {
        this.commandUtil.processExitWithError(err, JSON.stringify(result, undefined, 2));
      });
  }

  processFullPipelineInstance(fullPipeline: FullPipeline,
                              cbStatus: (err: Error, result: ProcessFullPipelineStatus) => void,
                              cbFinal: (err: Error, result?: ProcessFullPipelineResult) => void) {
    let me = this;
    cbStatus = me.checkCallback(cbStatus);
    cbFinal = me.checkCallback(cbFinal);
    try {
      const startTime = new Date();
      fullPipeline.validate();
      let processFullPipelineResult: ProcessFullPipelineResult = {};
      me.vitaDecryptUnTar.process(fullPipeline.decryptAndUnTarOptions,
        (err: Error, decryptAndUnTarStatus: DecryptAndUnTarStatus) => {
          cbStatus(err, {
            tag: fullPipeline.tag,
            decryptAndUnTarStatus,
            startTime,
            currentTime: new Date(),
            status: 'Decrypting ...'
          });
        },
        (err: Error, decryptAndUnTarResults: DecryptAndUnTarResult[]) => {
          processFullPipelineResult.decryptAndUnTarResults = decryptAndUnTarResults;
          if (me.commandUtil.callbackIfError(cbFinal, err, processFullPipelineResult)) {
            return;
          }
          decryptAndUnTarResults.forEach(result => {
            fullPipeline.unZipFileOptions.zippedFiles = result.zipFiles;
            processFullPipelineResult.unZipFileResults = [];
            me.vitaFileUtil.unZipFiles(fullPipeline.unZipFileOptions,
              (err: Error, unZipFileStatus: UnZipFileStatus) => {
                cbStatus(err, {
                  tag: fullPipeline.tag,
                  unZipFileStatus,
                  startTime,
                  currentTime: new Date(),
                  status: 'Unzipping PCAP file ...'
                });
              },
              (err: Error, unZipFileResults: UnZipFileResult[]) => {
                processFullPipelineResult.unZipFileResults.push(unZipFileResults);
                if (me.commandUtil.callbackIfError(cbFinal, err)) {
                  return;
                }
                fullPipeline.mergePcapFilesOptions.pcapFiles = unZipFileResults.map(unZipFileResult => {
                  return unZipFileResult.unzippedFilePath;
                }).filter(pcapFile => {
                  return !!pcapFile;
                });
                me.vitaPcapMerge.mergePcapFiles(fullPipeline.mergePcapFilesOptions,
                  (err: Error, mergePcapFilesStatus: MergePcapFilesStatus) => {
                    cbStatus(err, {
                      tag: fullPipeline.tag,
                      mergePcapFilesStatus,
                      startTime,
                      currentTime: new Date(),
                      status: 'Merging PCAP files ...'
                    });
                  },
                  (err: Error, result: MergePcapFilesResult) => {
                    processFullPipelineResult.mergePcapFilesResult = result;
                    processFullPipelineResult.tag = fullPipeline.tag;
                    let finalStatus: ProcessFullPipelineStatus = {
                      tag: fullPipeline.tag,
                      startTime,
                      currentTime: new Date(),
                      endTime: new Date(),
                      status: 'Finished.'
                    };
                    cbStatus(err, finalStatus);
                    deepExtend(processFullPipelineResult, finalStatus);
                    cbFinal(err, processFullPipelineResult);
                  });
              })
          });
        });
    } catch (err) {
      cbFinal(err);
    }
  }
}