Merge branch 'master' into bleepbop/VIRTS-2881/health-v2-api-pytest

mitre · Oct 6, 2021 · 86c3588 · 86c3588
2 parents 89c4112 + cc7e46e
commit 86c3588
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/app/api/v2/handlers/fact_api.py b/app/api/v2/handlers/fact_api.py
@@ -76,12 +76,12 @@ async def add_facts(self, request: web.Request):
             new_fact = Fact.load(fact_data)
             if 'source' not in fact_data:
                 new_fact.source = WILDCARD_STRING
-            new_fact.source_type = OriginType.USER.name
+            new_fact.origin_type = OriginType.USER
             await knowledge_svc_handle.add_fact(new_fact)
             store = await knowledge_svc_handle.get_facts(criteria=dict(trait=new_fact.trait,
                                                                        value=new_fact.value,
                                                                        source=WILDCARD_STRING,
-                                                                       source_type=OriginType.USER.name))
+                                                                       origin_type=OriginType.USER))
             resp = await self._api_manager.verify_fact_integrity(store)
             return web.json_response(dict(added=resp))
         except Exception as e:
@@ -105,10 +105,10 @@ async def add_relationships(self, request: web.Request):
             shorthand = new_relationship.shorthand
             new_relationship.source.relationships = [shorthand]
             new_relationship.source.source = origin_target
-            new_relationship.source.source_type = OriginType.USER.name
+            new_relationship.source.origin_type = OriginType.USER
             if 'target' in relationship_data:
                 new_relationship.target.source = origin_target
-                new_relationship.target.source_type = OriginType.USER.name
+                new_relationship.target.origin_type = OriginType.USER
                 new_relationship.target.relationships = [shorthand]
                 await knowledge_svc_handle.add_fact(new_relationship.target)
             await knowledge_svc_handle.add_fact(new_relationship.source)

diff --git a/app/objects/c_agent.py b/app/objects/c_agent.py
@@ -280,7 +280,7 @@ async def task(self, abilities, obfuscator, facts=(), deadman=False):
         knowledge_svc_handle = BaseService.get_service('knowledge_svc')
         for fact in facts:
             fact.source = self.paw
-            fact.origin_type = OriginType.SEEDED.name
+            fact.origin_type = OriginType.SEEDED
             await knowledge_svc_handle.add_fact(fact)
         self.links.extend(links)
         return links

diff --git a/app/service/learning_svc.py b/app/service/learning_svc.py
@@ -54,7 +54,7 @@ async def learn(self, facts, link, blob, operation=None):
 
     @staticmethod
     async def _save_fact(link, facts, fact, operation=None):
-        fact.source_type = OriginType.LEARNED.name
+        fact.origin_type = OriginType.LEARNED
         fact.source = operation.id if operation else link.id
         if all(fact.trait) and not any(fact == f for f in facts):
             fact.collected_by = link.paw