From 6866b49c87ac357b990fa77248f710394a997917 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= <raphael@vinot.info>
Date: Fri, 28 Aug 2020 18:24:14 +0200
Subject: [PATCH] chg: Add support for legitimate content context

---
 lookyloo/lookyloo.py    | 35 +++++++++++++++++++++++++++++++++++
 website/web/__init__.py |  2 +-
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/lookyloo/lookyloo.py b/lookyloo/lookyloo.py
index 5d2e52c9..06fde7af 100644
--- a/lookyloo/lookyloo.py
+++ b/lookyloo/lookyloo.py
@@ -187,6 +187,12 @@ class Context():
             elif filename == 'malicious':
                 for h, details in file_content.items():
                     p.sadd('bh|malicious', h)
+            elif filename == 'legitimate':
+                for h, details in file_content.items():
+                    if 'domain' in details:
+                        p.sadd(f'bh|{h}|legitimate', *details['domain'])
+                    elif 'description' in details:
+                        p.hset('known_content', h, details['description'])
             else:
                 for h, details in file_content.items():
                     p.sadd(f'bh|{h}|legitimate', *details['hostnames'])
@@ -323,6 +329,35 @@ class Context():
             p.sadd(f'{ressource_hash}|tag', details['type'])
         p.execute()
 
+    def store_known_legitimate_ressource(self, ressource_hash: str, details: Dict[str, str]):
+        known_legitimate_ressource_file = get_homedir() / 'known_content' / 'legitimate.json'
+        if known_legitimate_ressource_file.exists():
+            with open(known_legitimate_ressource_file) as f:
+                to_store = json.load(f)
+        else:
+            to_store = {}
+
+        if ressource_hash not in to_store:
+            to_store[ressource_hash] = {'domain': set(), 'description': ''}
+        else:
+            to_store[ressource_hash]['domain'] = set(to_store[ressource_hash]['domain'])
+
+        if 'domain' in details:
+            to_store[ressource_hash]['domain'].add(details['domain'])
+        if 'description' in details:
+            to_store[ressource_hash]['description'] = details['description']
+
+        with open(known_legitimate_ressource_file, 'w') as f:
+            json.dump(to_store, f, indent=2, default=dump_to_json)
+
+    def add_legitimate(self, ressource_hash: str, details: Dict[str, str]):
+        self.store_known_legitimate_ressource(ressource_hash, details)
+        if 'domain' in details:
+            self.redis.sadd(f'bh|{ressource_hash}|legitimate', details['domain'])
+        elif 'description' in details:
+            # Library
+            self.redis.hset('known_content', ressource_hash, details['description'])
+
     # Query DB
 
     def is_legitimate(self, urlnode: URLNode, known_hashes: Iterable[str]) -> Optional[bool]:
diff --git a/website/web/__init__.py b/website/web/__init__.py
index b98615ff..c441fb63 100644
--- a/website/web/__init__.py
+++ b/website/web/__init__.py
@@ -523,7 +523,7 @@ def add_context(tree_uuid: str, urlnode_uuid: str):
         if context_data.get('legitimate_domain'):
             legitimate_details['domain'] = context_data['legitimate_domain']
         if context_data.get('legitimate_description'):
-            legitimate_details['target'] = context_data['legitimate_description']
+            legitimate_details['description'] = context_data['legitimate_description']
         details['legitimate'] = legitimate_details
     lookyloo.add_context(tree_uuid, urlnode_uuid, ressource_hash, legitimate, malicious, details)
     return redirect(url_for('hostnode_popup', tree_uuid=tree_uuid, node_uuid=hostnode_uuid))