firmament-vita/ts/implementations/vita-pcap-merge-impl.ts

Firmament module for automating VITA tasks

import {injectable, inject} from 'inversify';
import {CommandUtil, ForceErrorImpl} from "firmament-yargs";
const fs = require('fs');
const async = require('async');
const path = require('path');
const mkdirp = require('mkdirp');
//-->
import {VitaFileUtil} from "../interfaces/vita-file-util";
import {
  ParallelMergePcapFilesResult, MergePcapFilesResult, MergePcapFilesOptions, MergePcapFilesStatus
} from "../interfaces/vita-options-results";
import {VitaSpawn} from "../interfaces/vita-spawn";
import {VitaPcapMerge} from "../interfaces/vita-pcap-merge";
//-->

@injectable()
export class VitaPcapMergeImpl extends ForceErrorImpl implements VitaPcapMerge {
  private useDummyMerge = false;
  private counter: number = 0;

  constructor(@inject('CommandUtil')private commandUtil: CommandUtil,
              @inject('VitaFileUtil')private vitaFileUtil: VitaFileUtil,
              @inject('VitaSpawn')private vitaSpawn: VitaSpawn) {
    super();
  }

  mergePcapFiles(mergePcapFileOptions: MergePcapFilesOptions,
                 cbStatus: (err: Error, mergePcapFilesStatus: MergePcapFilesStatus)=>void,
                 cbFinal: (err: Error, mergePcapFilesResults?: MergePcapFilesResult)=>void) {
    cbStatus = this.checkCallback(cbStatus);
    cbFinal = this.checkCallback(cbFinal);
    let mergedPcapFile = mergePcapFileOptions.mergedPcapFile;
    let pcapFiles = mergePcapFileOptions.pcapFiles;
    let totalBytes = VitaPcapMergeImpl.sumOfFileSizes(pcapFiles);
    let bytesWritten = 0;
    let args: string[] = ['-w', '-'];
    mkdirp(path.dirname(mergedPcapFile));
    let writeFileStream = fs.createWriteStream(mergedPcapFile);
    args = args.concat(pcapFiles);
    this.vitaSpawn.spawnIt(
      'mergecap', args,
      (err: Error, result: any) => {
        bytesWritten += result.length;
        let taskName = `Creating ${mergedPcapFile}`;
        let current = bytesWritten;
        let total = totalBytes;
        cbStatus(null, {taskName, current, total});
        writeFileStream.write(result);
      },
      (err: Error, result: any) => {
        writeFileStream.close();
        if (mergePcapFileOptions.deleteOriginalPcapFiles) {
          this.vitaFileUtil.deleteFiles(mergePcapFileOptions.pcapFiles);
        }
        cbFinal(null, {error: err, mergedPcapFile});
      }
    );
  }

  private static sumOfFileSizes(files: string[]): number {
    let retVal = 0;
    files.forEach(file => {
      retVal += VitaPcapMergeImpl.getFilesizeInBytes(file);
    });

    return retVal;
  }

  private static getFilesizeInBytes(filename) {
    let stats = fs.statSync(filename);
    let fileSizeInBytes = stats["size"];
    return fileSizeInBytes;
  }

  mergePcapFilesParallel(files: string[],
                         mergedFilename: string,
                         cb: (err: Error, mergePcapFilesResults?: ParallelMergePcapFilesResult[])=>void) {
    let me = this;
    cb = this.checkCallback(cb);
    files = files || [];
    let columns = [];
    let mask = 0x80000000;
    while (files.length) {
      while (mask && !(mask & files.length)) {
        mask >>>= 1;
      }
      columns.push(files.slice(0, mask));
      files = files.slice(mask);
    }
    let fnArray = [];
    for (let i = 0; i < columns.length; ++i) {
      let column = columns[i];
      fnArray.push(async.apply(this.reducePowerOf2FileListToOneFile.bind(this), column, me.useDummyMerge));
    }
    async.series(fnArray, (err: Error, results: ParallelMergePcapFilesResult[]) => {
      let me = this;
      files = results.map(result => result.outFile);
      if (files.length > 1) {
        this.mergePcapFilesParallel(files, mergedFilename, (err, res) => {
          cb(null, res);
        })
      } else {
        if (results.length !== 1) {
          cb(new Error('Bad results length'));
          return;
        }
        if (!me.useDummyMerge) {
          fs.rename(results[0].outFile, mergedFilename, function (err) {
            if (me.commandUtil.callbackIfError(cb, err)) {
              return;
            }
            fs.stat(mergedFilename, function (err, stats) {
              cb(null, [{
                error: err,
                inFile0: '',
                inFile1: '',
                outFile: mergedFilename,
                outFileStats: JSON.stringify(stats)
              }]);
            });
          });
        } else {
          cb(null, [{error: err, inFile0: '', inFile1: '', outFile: mergedFilename, outFileStats: 'Using DummyMerge'}]);
        }
      }
    });
  }

  private reducePowerOf2FileListToOneFile(files: string[], useDummyMerge: boolean, cb: (err: Error, result: ParallelMergePcapFilesResult)=>void) {
    let me = this;
    let triangleReduceRows = [];
    me.buildMergeTupleArray(files, 0, triangleReduceRows);
    let fnRowReduceSeriesArray = [];
    for (let j = 0; j < triangleReduceRows.length; ++j) {
      let triangleReduceRow = triangleReduceRows[j];
      let fnRowReduceParallelArray = [];
      triangleReduceRow.forEach(twoToOneReduce => {
        let inFile0 = twoToOneReduce.inFile0;
        let inFile1 = twoToOneReduce.inFile1;
        let outFile = twoToOneReduce.outFile;
        if (!useDummyMerge) {
          if (!inFile1) {
            fnRowReduceParallelArray.push(async.apply(me.vitaFileUtil.copyFile.bind(me.vitaFileUtil), inFile0, outFile));
          } else {
            fnRowReduceParallelArray.push(async.apply(me.spawnMergeTwoPcapsOperation.bind(me), inFile0, inFile1, outFile));
          }
        } else {
          fnRowReduceParallelArray.push(async.apply(VitaPcapMergeImpl.dummyMergeOperation.bind(me), inFile0, inFile1, outFile));
        }
      });
      fnRowReduceSeriesArray.push(async.apply(async.parallelLimit, fnRowReduceParallelArray, 3));
    }
    async.series(fnRowReduceSeriesArray, (err: Error, resultAryAry: ParallelMergePcapFilesResult[][]) => {
      let resultAry = resultAryAry[resultAryAry.length - 1];
      let retVal = resultAry[0];
      let filesToDelete: string[] = [];
      while (resultAry = resultAryAry.pop()) {
        resultAry.forEach(result => {
          if (VitaPcapMergeImpl.isIntermediateFile(result.inFile0)) {
            filesToDelete.push(result.inFile0);
          }
          if (VitaPcapMergeImpl.isIntermediateFile(result.inFile1)) {
            filesToDelete.push(result.inFile1);
          }
        });
      }
      this.vitaFileUtil.deleteFiles(filesToDelete);
      cb(err, retVal);
    });
  }

  private buildMergeTupleArray(files: string[], depth: number, parallelTuples: any[]) {
    let tuples = [];
    for (let i = 0; i < files.length; i += 2) {
      let inFile0 = files[i];
      let inFile1 = files[i + 1];
      let workingDir = path.dirname(inFile0);
      let baseFile = path.basename(inFile0);
      let outFile = this.getIntermediateFileName(baseFile);
      outFile = path.resolve(workingDir, outFile);
      tuples.push({inFile0, inFile1, outFile});
    }
    parallelTuples.push(tuples);
    files = tuples.map(tuple => tuple.outFile);
    if (files.length > 1) {
      this.buildMergeTupleArray(files, depth + 1, parallelTuples);
    }
  }

  private static isIntermediateFile(fileName): boolean {
    return (/\d{6}/.test(fileName) && /\.tmp$/.test(fileName));
  }

  private getIntermediateFileName(fileName): string {
    let retVal: string;
    let prefix = VitaPcapMergeImpl.zeroFill(this.counter++, 6);
    //does basename start with 6 digits and end with '.tmp'?
    if (VitaPcapMergeImpl.isIntermediateFile(fileName)) {
      //strip 6 digits from front and replace with new 6 digits
      retVal = `${prefix}${fileName.substr(6)}`;
    } else {
      //Pre and post fix our identifiers for which files to delete
      retVal = `${prefix}_${fileName}.tmp`;
    }
    this.commandUtil.log(retVal);
    return retVal;
  }

  private static dummyMergeOperation(inFile0: string,
                                     inFile1: string,
                                     outFile: string,
                                     cb: (err: Error, mergePcapFilesResult: ParallelMergePcapFilesResult)=>void) {
    cb(null, {error: null, inFile0, inFile1, outFile, outFileStats: 'Using DummyMerge'});
  }

  private spawnMergeTwoPcapsOperation(inFile0: string,
                                      inFile1: string,
                                      outFile: string,
                                      cb: (err: Error, mergePcapFilesResult: ParallelMergePcapFilesResult)=>void) {
    this.vitaSpawn.spawnIt(
      'mergecap', [
        '-w',
        outFile,
        inFile0,
        inFile1
      ],
      (err: Error, result: any) => {
        // cb(null, null);
      },
      (err: Error, result: any) => {
        cb(null, {error: err, inFile0, inFile1, outFile, outFileStats: 'Nobody Cares (intermediate file)'});
      }
    );
  }

  static zeroFill(number, width) {
    width -= number.toString().length;
    if (width > 0) {
      return new Array(width + (/\./.test(number) ? 2 : 1)).join('0') + number;
    }
    return number + ""; // always return a string
  }
}