chg: Fix typing

lookyloo/exceptions.py

@@ -20,3 +20,7 @@ class CreateDirectoryException(LookylooException):
 
 class ConfigError(LookylooException):
     pass
+
+
+class MissingUUID(LookylooException):
+    pass

lookyloo/lookyloo.py

@@ -22,7 +22,7 @@ from har2tree import CrawledTree, Har2TreeError, HarFile, HostNode, URLNode
 from redis import Redis
 from scrapysplashwrapper import crawl
 
-from .exceptions import NoValidHarFile
+from .exceptions import NoValidHarFile, MissingUUID
 from .helpers import get_homedir, get_socket_path, load_cookies, load_configs, safe_create_dir, get_email_template
 from .modules import VirusTotal, SaneJavaScript
 
@@ -86,10 +86,14 @@ class Lookyloo():
 
     def get_urlnode_from_tree(self, capture_dir: Path, node_uuid: str) -> URLNode:
         ct = self._load_pickle(capture_dir / 'tree.pickle')
+        if not ct:
+            raise MissingUUID(f'Unable to find UUID {node_uuid} in {capture_dir}')
         return ct.root_hartree.get_url_node_by_uuid(node_uuid)
 
     def get_hostnode_from_tree(self, capture_dir: Path, node_uuid: str) -> HostNode:
         ct = self._load_pickle(capture_dir / 'tree.pickle')
+        if not ct:
+            raise MissingUUID(f'Unable to find UUID {node_uuid} in {capture_dir}')
         return ct.root_hartree.get_host_node_by_uuid(node_uuid)
 
     def get_statistics(self, capture_dir: Path) -> Dict[str, Any]:
@@ -272,7 +276,7 @@ class Lookyloo():
         except Exception as e:
             logging.exception(e)
 
-    def load_tree(self, capture_dir: Path) -> Tuple[str, str, str, str, str, Dict[str, str]]:
+    def load_tree(self, capture_dir: Path) -> Tuple[str, str, str, str, Dict[str, str]]:
         har_files = sorted(capture_dir.glob('*.har'))
         pickle_file = capture_dir / 'tree.pickle'
         try:

lookyloo/modules.py

@@ -35,7 +35,7 @@ class SaneJavaScript():
         self.storage_dir = get_homedir() / 'sanejs'
         self.storage_dir.mkdir(parents=True, exist_ok=True)
 
-    def hashes_lookup(self, sha512: Union[List[str], str], force: bool=False) -> Optional[Dict[str, Any]]:
+    def hashes_lookup(self, sha512: Union[List[str], str], force: bool=False) -> Dict[str, Any]:
         if isinstance(sha512, str):
             hashes = [sha512]
         else:

poetry.lock

@@ -870,7 +870,7 @@ scrapy = "^1.8.0"
 scrapy-splash = "^0.7.2"
 
 [package.source]
-reference = "a67b329d4b4281a90418c2e63464523294af6d53"
+reference = "2713eaf0808eaa1244cdf0df96db20138670d16d"
 type = "git"
 url = "https://github.com/viper-framework/ScrapySplashWrapper.git"
 [[package]]

website/web/__init__.py

@@ -109,6 +109,8 @@ def scrape_web():
 @app.route('/tree/<string:tree_uuid>/hostname/<string:node_uuid>/text', methods=['GET'])
 def hostnode_details_text(tree_uuid: str, node_uuid: str):
     capture_dir = lookyloo.lookup_capture_dir(tree_uuid)
+    if not capture_dir:
+        return
     hostnode = lookyloo.get_hostnode_from_tree(capture_dir, node_uuid)
     urls = []
     for url in hostnode.urls:
@@ -126,6 +128,8 @@ def hostnode_details_text(tree_uuid: str, node_uuid: str):
 @app.route('/tree/<string:tree_uuid>/hostname_popup/<string:node_uuid>', methods=['GET'])
 def hostnode_popup(tree_uuid: str, node_uuid: str):
     capture_dir = lookyloo.lookup_capture_dir(tree_uuid)
+    if not capture_dir:
+        return
     hostnode = lookyloo.get_hostnode_from_tree(capture_dir, node_uuid)
     table_keys = {
         'js': "/static/javascript.png",
@@ -165,6 +169,8 @@ def hostnode_popup(tree_uuid: str, node_uuid: str):
 @app.route('/tree/<string:tree_uuid>/url/<string:node_uuid>', methods=['GET'])
 def urlnode_details(tree_uuid: str, node_uuid: str):
     capture_dir = lookyloo.lookup_capture_dir(tree_uuid)
+    if not capture_dir:
+        return
     urlnode = lookyloo.get_urlnode_from_tree(capture_dir, node_uuid)
     to_return = BytesIO()
     got_content = False