Just some refactoring

and chages to pass the build
This commit is contained in:
Shreya 2020-08-23 23:17:58 +05:30
parent 2431d45b74
commit e25e913e86
2 changed files with 22 additions and 19 deletions

View File

@ -1,9 +1,9 @@
import subprocess import subprocess
from common.data.post_breach_consts import POST_BREACH_CLEAR_CMD_HISTORY from common.data.post_breach_consts import POST_BREACH_CLEAR_CMD_HISTORY
from infection_monkey.post_breach.pba import PBA
from infection_monkey.post_breach.clear_command_history.clear_command_history import \ from infection_monkey.post_breach.clear_command_history.clear_command_history import \
get_commands_to_clear_command_history get_commands_to_clear_command_history
from infection_monkey.post_breach.pba import PBA
from infection_monkey.telemetry.post_breach_telem import PostBreachTelem from infection_monkey.telemetry.post_breach_telem import PostBreachTelem

View File

@ -6,21 +6,18 @@ export default function parsePbaResults(results) {
const SHELL_STARTUP_NAME = 'Modify shell startup file'; const SHELL_STARTUP_NAME = 'Modify shell startup file';
const CMD_HISTORY_NAME = 'Clear command history'; const CMD_HISTORY_NAME = 'Clear command history';
const multipleResultsPbas = [SHELL_STARTUP_NAME, CMD_HISTORY_NAME]
function aggregateMultipleResultsPba(results) { function aggregateMultipleResultsPba(results) {
let aggregatedPbaResults = { let aggregatedPbaResults = {};
'Modify shell startup file': { multipleResultsPbas.forEach(function(pba) {
aggregatedResult: undefined, aggregatedPbaResults[pba] = {
successfulOutputs: '',
failedOutputs: '',
isSuccess: false
},
'Clear command history': {
aggregatedResult: undefined, aggregatedResult: undefined,
successfulOutputs: '', successfulOutputs: '',
failedOutputs: '', failedOutputs: '',
isSuccess: false isSuccess: false
} }
} })
function aggregateResults(result) { function aggregateResults(result) {
if (aggregatedPbaResults[result.name].aggregatedResult === undefined) { if (aggregatedPbaResults[result.name].aggregatedResult === undefined) {
@ -30,30 +27,36 @@ function aggregateMultipleResultsPba(results) {
aggregatedPbaResults[result.name].successfulOutputs += result.result[0]; aggregatedPbaResults[result.name].successfulOutputs += result.result[0];
aggregatedPbaResults[result.name].isSuccess = true; aggregatedPbaResults[result.name].isSuccess = true;
} }
if (!result.result[1]) { else if (!result.result[1]) {
aggregatedPbaResults[result.name].failedOutputs += result.result[0]; aggregatedPbaResults[result.name].failedOutputs += result.result[0];
} }
} }
function checkAggregatedResults(pbaName) { function checkAggregatedResults(pbaName) { // if this pba's results were aggregated, push to `results`
if (aggregatedPbaResults[pbaName].aggregatedResult !== undefined) { if (aggregatedPbaResults[pbaName].aggregatedResult !== undefined) {
aggregatedPbaResults[pbaName].aggregatedResult.result[0] = aggregatedPbaResults[pbaName].successfulOutputs + aggregatedPbaResults[pbaName].failedOutputs; aggregatedPbaResults[pbaName].aggregatedResult.result[0] = (aggregatedPbaResults[pbaName].successfulOutputs +
aggregatedPbaResults[pbaName].failedOutputs);
aggregatedPbaResults[pbaName].aggregatedResult.result[1] = aggregatedPbaResults[pbaName].isSuccess; aggregatedPbaResults[pbaName].aggregatedResult.result[1] = aggregatedPbaResults[pbaName].isSuccess;
results.push(aggregatedPbaResults[pbaName].aggregatedResult); results.push(aggregatedPbaResults[pbaName].aggregatedResult);
} }
} }
// check for pbas with multiple results and aggregate their results
for (let i = 0; i < results.length; i++) for (let i = 0; i < results.length; i++)
if (results[i].name === SHELL_STARTUP_NAME || results[i].name === CMD_HISTORY_NAME) if (multipleResultsPbas.includes(results[i].name))
aggregateResults(results[i]); aggregateResults(results[i]);
if (aggregatedPbaResults[SHELL_STARTUP_NAME].aggregatedResult === undefined && // if no modifications were made to the results, i.e. if no pbas had mutiple results, return `results` as it is
aggregatedPbaResults[CMD_HISTORY_NAME].aggregatedResult === undefined) let noResultsModifications = true;
multipleResultsPbas.forEach((pba) => {
if (aggregatedPbaResults[pba].aggregatedResult !== undefined)
noResultsModifications = false;
})
if (noResultsModifications)
return results; return results;
// if modifications were made, push aggregated results to `results` and return
results = results.filter(result => result.name !== SHELL_STARTUP_NAME && result.name !== CMD_HISTORY_NAME); results = results.filter(result => result.name !== SHELL_STARTUP_NAME && result.name !== CMD_HISTORY_NAME);
checkAggregatedResults(SHELL_STARTUP_NAME); multipleResultsPbas.forEach(pba => checkAggregatedResults(pba));
checkAggregatedResults(CMD_HISTORY_NAME);
return results; return results;
} }