attack link to org and feedback save task - first pass

Author: mkaranasou

src/baskerville/db/models.py

@@ -212,6 +212,7 @@ class Attack(Base, SerializableMixin):
 
     id = Column(BigInteger, primary_key=True)
     id_misp = Column(BigInteger)
+    uuid_org = Column(TEXT())
     date = Column(DateTime(timezone=True))
     start = Column(DateTime(timezone=True))
     stop = Column(DateTime(timezone=True))
@@ -224,6 +225,7 @@ class Attack(Base, SerializableMixin):
     sync_stop = Column(DateTime(timezone=True))
     processed = Column(Integer)
     notes = Column(TEXT)
+    progress_report = Column(TEXT)
     analysis_notebook = Column(TEXT)
 
     request_sets = relationship(
@@ -234,6 +236,10 @@ class Attack(Base, SerializableMixin):
         'Attribute', secondary='attribute_attack_link',
         back_populates='attacks'
     )
+    organization = relationship(
+        'Organization',
+        primaryjoin='foreign(Attack.uuid_org) == remote(Organization.uuid)'
+    )
 
 
 class Attribute(Base, SerializableMixin):

src/baskerville/models/pipeline_tasks/feedback_pipeline.py

@@ -6,7 +6,7 @@
 from baskerville.db.dashboard_models import Feedback
 from baskerville.models.pipeline_tasks.tasks_base import Task
 from baskerville.models.config import BaskervilleConfig
-from baskerville.models.pipeline_tasks.tasks import GetDataKafka, Save
+from baskerville.models.pipeline_tasks.tasks import GetDataKafka, SaveFeedback
 
 
 def set_up_feedback_pipeline(config: BaskervilleConfig):
@@ -21,10 +21,10 @@ def set_up_feedback_pipeline(config: BaskervilleConfig):
         GetDataKafka(
             config,
             steps=[
-                Save(
+                SaveFeedback(
                     config,
                     table_model=Feedback,
-                    not_common=()
+                    not_common=('feedback_context', 'progress_count')
                 ),
             ]),
     ]

src/baskerville/models/pipeline_tasks/tasks.py

@@ -20,9 +20,10 @@
 from pyspark.streaming import StreamingContext
 from functools import reduce
 from pyspark.sql import DataFrame
+from sqlalchemy.exc import SQLAlchemyError
 
 from baskerville.db import get_jdbc_url
-from baskerville.db.models import RequestSet, Model
+from baskerville.db.models import RequestSet, Model, Attack
 from baskerville.models.banjax_report_consumer import BanjaxReportConsumer
 from baskerville.models.ip_cache import IPCache
 from baskerville.models.metrics.registry import metrics_registry
@@ -933,6 +934,50 @@ def run(self):
         return self.df
 
 
+class SaveFeedback(Save):
+    def upsert_attack(self):
+        new_ = False
+        success = False
+        try:
+            feedback_context = self.df.select('feedback_context').collect()
+            attack = self.db_tools.session.query(Attack).filter_by(
+                    uuid=feedback_context.uuid
+            ).filter_by(
+                start=feedback_context.start
+            ).filter_by(stop=feedback_context.stop).first()
+            if not attack:
+                attack = Attack()
+                new_ = True
+            attack.uuid_org = feedback_context.uuid
+            attack.start = feedback_context.start
+            attack.stop = feedback_context.stop
+            attack.target = feedback_context.target
+            attack.ip_count = feedback_context.ip_count
+            attack.notes = feedback_context.notes
+            attack.progress_report = feedback_context.progress_report
+            if new_:
+                self.db_tools.session.add(attack)
+            self.db_tools.session.commit()
+            success = True
+        except SQLAlchemyError as sqle:
+            traceback.print_exc()
+            self.db_tools.session.rollback()
+            success = False
+            self.logger.error(str(sqle))
+            # todo: what should the handling be?
+        except Exception as e:
+            traceback.print_exc()
+            success = False
+            self.logger.error(str(e))
+            # todo: what should the handling be?
+        return success
+
+    def prepare_to_save(self):
+        success = self.upsert_attack()
+        if success:
+            super(SaveFeedback, self).prepare_to_save()
+
+
 class RefreshCache(CacheTask):
     def run(self):
         self.service_provider.refresh_cache(self.df)