From cc1f46af32c6548dbc28330fc916602c8bfbd713 Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Fri, 22 Nov 2019 18:30:08 +0200
Subject: [PATCH] Moved attack matrix to attack report page

---
 .../cc/resources/attack/attack_report.py      | 10 ++-
 .../cc/services/attack/attack_config.py       | 12 +++
 .../cc/services/attack/attack_report.py       | 15 ++--
 .../report-components/AttackReport.js         | 27 +++---
 .../report-components/attack/Matrix.js        | 82 +++++++++++++++++++
 .../src/components/ui-components/Checkbox.js  | 12 ++-
 6 files changed, 138 insertions(+), 20 deletions(-)
 create mode 100644 monkey/monkey_island/cc/ui/src/components/report-components/attack/Matrix.js

diff --git a/monkey/monkey_island/cc/resources/attack/attack_report.py b/monkey/monkey_island/cc/resources/attack/attack_report.py
index a55137d7e..dc63f6747 100644
--- a/monkey/monkey_island/cc/resources/attack/attack_report.py
+++ b/monkey/monkey_island/cc/resources/attack/attack_report.py
@@ -1,7 +1,8 @@
 import flask_restful
-from flask import jsonify
 from monkey_island.cc.auth import jwt_required
 from monkey_island.cc.services.attack.attack_report import AttackReportService
+from monkey_island.cc.services.attack.attack_schema import SCHEMA
+from flask import json, current_app
 
 __author__ = "VakarisZ"
 
@@ -10,4 +11,9 @@ class AttackReport(flask_restful.Resource):
 
     @jwt_required()
     def get(self):
-        return jsonify(AttackReportService.get_latest_report()['techniques'])
+        response_content = {'techniques': AttackReportService.get_latest_report()['techniques'], 'schema': SCHEMA}
+        return current_app.response_class(json.dumps(response_content,
+                                                     indent=None,
+                                                     separators=(",", ":"),
+                                                     sort_keys=False) + "\n",
+                                          mimetype=current_app.config['JSONIFY_MIMETYPE'])
diff --git a/monkey/monkey_island/cc/services/attack/attack_config.py b/monkey/monkey_island/cc/services/attack/attack_config.py
index 6feff990e..f834f259b 100644
--- a/monkey/monkey_island/cc/services/attack/attack_config.py
+++ b/monkey/monkey_island/cc/services/attack/attack_config.py
@@ -173,3 +173,15 @@ class AttackConfig(object):
             for key, technique in list(attack_type['properties'].items()):
                 techniques[key] = technique['value']
         return techniques
+
+    @staticmethod
+    def get_techniques_for_report():
+        """
+        :return: Format: {"T1110": {"selected": True, "type": "Credential Access", "T1075": ...}
+        """
+        attack_config = AttackConfig.get_config()
+        techniques = {}
+        for type_name, attack_type in list(attack_config.items()):
+            for key, technique in list(attack_type['properties'].items()):
+                    techniques[key] = {'selected': technique['value'], 'type': SCHEMA['properties'][type_name]['title']}
+        return techniques
diff --git a/monkey/monkey_island/cc/services/attack/attack_report.py b/monkey/monkey_island/cc/services/attack/attack_report.py
index 10005bd26..43cca0876 100644
--- a/monkey/monkey_island/cc/services/attack/attack_report.py
+++ b/monkey/monkey_island/cc/services/attack/attack_report.py
@@ -58,13 +58,14 @@ class AttackReportService:
                 'name': REPORT_NAME
             }
 
-        for tech_id, value in list(AttackConfig.get_technique_values().items()):
-            if value:
-                try:
-                    report['techniques'].update({tech_id: TECHNIQUES[tech_id].get_report_data()})
-                except KeyError as e:
-                    LOG.error("Attack technique does not have it's report component added "
-                              "to attack report service. %s" % e)
+        for tech_id, tech_info in list(AttackConfig.get_techniques_for_report().items()):
+            try:
+                technique_report_data = TECHNIQUES[tech_id].get_report_data()
+                technique_report_data.update(tech_info)
+                report['techniques'].update({tech_id: technique_report_data})
+            except KeyError as e:
+                LOG.error("Attack technique does not have it's report component added "
+                          "to attack report service. %s" % e)
         mongo.db.attack_report.replace_one({'name': REPORT_NAME}, report, upsert=True)
         return report
 
diff --git a/monkey/monkey_island/cc/ui/src/components/report-components/AttackReport.js b/monkey/monkey_island/cc/ui/src/components/report-components/AttackReport.js
index add32ad4a..7c5f6cdc4 100644
--- a/monkey/monkey_island/cc/ui/src/components/report-components/AttackReport.js
+++ b/monkey/monkey_island/cc/ui/src/components/report-components/AttackReport.js
@@ -6,6 +6,7 @@ import '../../styles/Collapse.scss';
 import AuthComponent from '../AuthComponent';
 import {ScanStatus} from '../attack/techniques/Helpers';
 import Collapse from '@kunukn/react-collapse';
+import Matrix from './attack/Matrix';
 
 import T1210 from '../attack/techniques/T1210';
 import T1197 from '../attack/techniques/T1197';
@@ -66,7 +67,8 @@ class AttackReportPageComponent extends AuthComponent {
   constructor(props) {
     super(props);
     this.state = {
-      report: this.props.report,
+      techniques: this.props.report['techniques'],
+      schema: this.props.report['schema'],
       collapseOpen: ''
     };
   }
@@ -81,7 +83,7 @@ class AttackReportPageComponent extends AuthComponent {
     this.setState(state => ({collapseOpen: state.collapseOpen === technique ? null : technique}));
 
   getComponentClass(tech_id) {
-    switch (this.state.report[tech_id].status) {
+    switch (this.state.techniques[tech_id].status) {
       case ScanStatus.SCANNED:
         return 'collapse-info';
       case ScanStatus.USED:
@@ -95,7 +97,7 @@ class AttackReportPageComponent extends AuthComponent {
     return (
       <div key={tech_id} className={classNames('collapse-item', {'item--active': this.state.collapseOpen === tech_id})}>
         <button className={classNames('btn-collapse', this.getComponentClass(tech_id))} onClick={() => this.onToggle(tech_id)}>
-          <span>{this.state.report[tech_id].title}</span>
+          <span>{this.state.techniques[tech_id].title}</span>
           <span>
               <i className={classNames('fa', this.state.collapseOpen === tech_id ? 'fa-chevron-down' : 'fa-chevron-up')}></i>
           </span>
@@ -118,22 +120,26 @@ class AttackReportPageComponent extends AuthComponent {
     const TechniqueComponent = tech_components[technique];
     return (
       <div className={`content ${collapseState}`}>
-        <TechniqueComponent data={this.state.report[technique]} reportData={this.props.reportData}/>
+        <TechniqueComponent data={this.state.techniques[technique]}/>
       </div>
     );
   }
 
   renderLegend() {
     return (<div id="header" className="row justify-content-between attack-legend">
-      <Col xs={4}>
+      <Col xs={3}>
+          <i className="fa fa-circle-thin icon-unchecked"></i>
+          <span> - Dissabled</span>
+      </Col>
+      <Col xs={3}>
         <i className="fa fa-circle icon-default"></i>
         <span> - Unscanned</span>
       </Col>
-      <Col xs={4}>
+      <Col xs={3}>
         <i className="fa fa-circle icon-info"></i>
         <span> - Scanned</span>
       </Col>
-      <Col xs={4}>
+      <Col xs={3}>
         <i className="fa fa-circle icon-danger"></i>
         <span> - Used</span>
       </Col>
@@ -142,7 +148,7 @@ class AttackReportPageComponent extends AuthComponent {
 
   generateReportContent() {
     let content = [];
-    Object.keys(this.state.report).forEach((tech_id) => {
+    Object.keys(this.state.techniques).forEach((tech_id) => {
       content.push(this.getTechniqueCollapse(tech_id))
     });
     return (
@@ -153,8 +159,9 @@ class AttackReportPageComponent extends AuthComponent {
       <p>
         This report shows information about ATT&CK techniques used by Infection Monkey.
       </p>
+      {this.renderLegend()}
+      <Matrix techniques={this.state.techniques} schema={this.state.schema}/>
       <div>
-        {this.renderLegend()}
         <section className="attack-report">{content}</section>
       </div>
       <br/>
@@ -163,7 +170,7 @@ class AttackReportPageComponent extends AuthComponent {
   }
 
   render() {
-    if (Object.keys(this.state.report).length === 0 && this.state.runStarted) {
+    if (Object.keys(this.state.techniques).length === 0 && this.state.runStarted) {
         return (<h1>No techniques were scanned</h1>);
     } else {
       return (<div> {this.generateReportContent()}</div>);
diff --git a/monkey/monkey_island/cc/ui/src/components/report-components/attack/Matrix.js b/monkey/monkey_island/cc/ui/src/components/report-components/attack/Matrix.js
new file mode 100644
index 000000000..f37fbe9c9
--- /dev/null
+++ b/monkey/monkey_island/cc/ui/src/components/report-components/attack/Matrix.js
@@ -0,0 +1,82 @@
+import React from 'react';
+import Checkbox from '../../ui-components/Checkbox';
+import ReactTable from 'react-table';
+import 'filepond/dist/filepond.min.css';
+
+class MatrixComponent extends React.Component {
+  constructor(props) {
+    super(props);
+    this.state = {techniques: this.props.techniques,
+                  schema: this.props.schema}
+  }
+
+  getColumns() {
+    let columns = [];
+    for(let type_key in this.state.schema.properties){
+      if (! this.state.schema.properties.hasOwnProperty(type_key)){
+        continue;
+      }
+      let tech_type = this.state.schema.properties[type_key];
+      columns.push({
+        Header: tech_type.title,
+        id: type_key,
+        accessor: x => MatrixComponent.renderTechnique(x[tech_type.title]),
+        style: {'whiteSpace': 'unset'}
+      });
+    }
+    return columns;
+  }
+
+  getTableRows() {
+    let rows = [];
+    for (let tech_id in this.state.techniques) {
+      if (this.state.techniques.hasOwnProperty(tech_id)){
+        let technique_added = false;
+        let technique = this.state.techniques[tech_id];
+        for(let row of rows){
+          if (! row.hasOwnProperty(technique.type)){
+            row[technique.type] = technique;
+            technique_added = true;
+            break;
+          }
+        }
+        if (! technique_added){
+          let newRow = {};
+          newRow[technique.type] = technique;
+          rows.push(newRow)
+        }
+      }
+    }
+    return rows;
+  }
+
+  static renderTechnique(technique) {
+    if (technique == null || typeof technique === undefined) {
+      return (<div/>)
+    } else {
+      return (
+        <Checkbox checked={technique.selected}
+                  necessary={true}
+                  name={technique.title}
+                  changeHandler={function(){}}
+                  status={technique.status}>
+          {technique.title}
+        </Checkbox>)
+    }
+  }
+
+  render() {
+    let tableRows = this.getTableRows();
+    return (
+      <div>
+        <div className={'attack-matrix'}>
+          <ReactTable columns={this.getColumns()}
+                      data={tableRows}
+                      showPagination={false}
+                      defaultPageSize={tableRows.length}/>
+        </div>
+      </div>);
+  }
+}
+
+export default MatrixComponent;
diff --git a/monkey/monkey_island/cc/ui/src/components/ui-components/Checkbox.js b/monkey/monkey_island/cc/ui/src/components/ui-components/Checkbox.js
index e143e0aba..ebd40aa3d 100644
--- a/monkey/monkey_island/cc/ui/src/components/ui-components/Checkbox.js
+++ b/monkey/monkey_island/cc/ui/src/components/ui-components/Checkbox.js
@@ -14,10 +14,18 @@ class CheckboxComponent extends React.PureComponent {
 	changeHandler(name, checked) function will be called with these parameters:
 	this.props.name (the name of this component) and
 	this.state.checked (boolean indicating if this component is checked or not)
+	this.props.status (int) adds a class "status-x" to this checkbox. Used for styling.
   */
   constructor(props) {
     super(props);
+    if (this.props.hasOwnProperty("status")){
+      this.status = this.props.status;
+    } else {
+      this.status = false
+    }
+    console.log(this.props);
     this.state = {
+      status: this.status,
       checked: this.props.checked,
       necessary: this.props.necessary,
       isAnimating: false
@@ -48,6 +56,9 @@ class CheckboxComponent extends React.PureComponent {
   // Creates class string for component
   composeStateClasses(core) {
     let result = core;
+    if (this.state.status) {
+      result += ' status-'+this.state.status;
+    }
     if (this.state.necessary) {
       return result + ' blocked'
     }
@@ -56,7 +67,6 @@ class CheckboxComponent extends React.PureComponent {
     } else {
       result += ' is-unchecked'
     }
-
     if (this.state.isAnimating) {
       result += ' do-ping';
     }