Improve exception handling

maldump/__main__.py

@@ -5,7 +5,9 @@
 import argparse
 import csv
 import ctypes
+import getpass
 import io
+import logging
 import os
 import sys
 import tarfile
@@ -20,14 +22,19 @@
     from maldump.structures import Quarantine
 
 __version__ = "0.5.0"
+logger = logging.getLogger(__name__)
 
 
 def main() -> None:
     init()
     args = parse_cli()
+    init_logging(args.log_level)
 
     # Admin privileges are required for optimal function (windows only)
     if sys.platform == "win32" and not ctypes.windll.shell32.IsUserAnAdmin():
+        logger.critical(
+            "The program executed on Windows machine without proper privileges"
+        )
         print("Please try again with admin privileges")
         sys.exit(1)
 
@@ -37,9 +44,15 @@ def main() -> None:
     # Switch to root partition
     os.chdir(args.root_dir)
 
+    logger.debug(
+        'Working in directory "%s", files would be stored into "%s"', os.getcwd(), dest
+    )
+
     # Get a list of all installed avs
     avs = AVManager.detect()
 
+    logger.debug("Detected AVs: %s", [av.name for av in avs])
+
     if args.quar:
         export_files(avs, dest)
     elif args.meta:
@@ -156,6 +169,13 @@ def parse_cli() -> argparse.Namespace:
     parser.add_argument(
         "-a", "--all", action="store_true", help="equivalent of running both -q and -m"
     )
+    parser.add_argument(
+        "-t",
+        "--log-level",
+        choices=["critical", "fatal", "error", "warn", "warning", "info", "debug"],
+        default="warning",
+        help="log level",
+    )
     parser.add_argument(
         "-v", "--version", action="version", version="%(prog)s " + __version__
     )
@@ -170,5 +190,21 @@ def parse_cli() -> argparse.Namespace:
     return parser.parse_args()
 
 
+def init_logging(log_level: str) -> None:
+    numeric_level = getattr(logging, log_level.upper(), None)
+    if not isinstance(numeric_level, int):
+        raise ValueError("Invalid log level: " + log_level)  # noqa: TRY004
+    logging.basicConfig(
+        handlers=[
+            # logging.FileHandler("syslog.log", mode="w", encoding="utf-8"),
+            logging.StreamHandler(sys.stderr)
+        ],
+        level=numeric_level,
+        format="%(asctime)s:%(levelname)s:%(name)s:%(module)s:%(message)s",
+    )
+    logger.debug("Logging started, logger initialized successfully")
+    logger.info("Logging as user %s", getpass.getuser())
+
+
 if __name__ == "__main__":
     main()

maldump/constants.py

@@ -1,4 +1,6 @@
 from enum import Enum
+from typing import Any
+from xml.etree.ElementTree import Element
 
 
 class OperatingSystem(Enum):
@@ -9,3 +11,35 @@ class OperatingSystem(Enum):
 
 class ThreatMetadata(str, Enum):
     UNKNOWN_THREAT = "Unknown-no-metadata"
+
+
+class UnloggedObjects:
+    @staticmethod
+    def __contains__(item: Any) -> bool:
+        from maldump.parsers.avast_parser import AvastParser
+        from maldump.parsers.avg_parser import AVGParser
+        from maldump.parsers.eset_parser import EsetParser
+        from maldump.parsers.forticlient_parser import ForticlientParser
+        from maldump.parsers.kaitai.forticlient_parser import (
+            ForticlientParser as ForticlientKaitaiParser,
+        )
+        from maldump.parsers.kaspersky_parser import KasperskyParser
+        from maldump.parsers.malwarebytes_parser import MalwarebytesParser
+        from maldump.parsers.mcafee_parser import McafeeParser
+        from maldump.parsers.windef_parser import WindowsDefenderParser
+
+        unlogged = {
+            bytes,
+            EsetParser,
+            AvastParser,
+            AVGParser,
+            ForticlientParser,
+            KasperskyParser,
+            MalwarebytesParser,
+            McafeeParser,
+            WindowsDefenderParser,
+            ForticlientKaitaiParser.Timestamp,
+            Element,
+        }
+
+        return item in unlogged

maldump/parsers/avira_parser.py

@@ -1,25 +1,36 @@
 from __future__ import annotations
 
-from datetime import datetime as dt
+import logging
 
 from maldump.parsers.kaitai.avira_parser import AviraParser as KaitaiParser
 from maldump.structures import Parser, QuarEntry
+from maldump.utils import Parser as parse
+
+logger = logging.getLogger(__name__)
 
 
 class AviraParser(Parser):
-    def parse_from_log(self, _=None) -> dict[str, QuarEntry]:
+    def parse_from_log(self, data=None):
+        pass
+
+    def parse_from_fs(self, _=None) -> dict[str, QuarEntry]:
+        logger.info("Parsing from filesystem in %s", self.name)
         quarfiles = {}
-        for metafile in self.location.glob("*.qua"):
-            kt = KaitaiParser.from_file(metafile)
+        for idx, metafile in enumerate(self.location.glob("*.qua")):
+            logger.debug('Parsing entry, idx %s, path "%s"', idx, metafile)
+
+            kt = parse(self).kaitai(KaitaiParser, metafile)
+            if kt is None:
+                logger.debug('Skipping entry idx %s, path "%s"', idx, metafile)
+                continue
+
             q = QuarEntry()
-            q.timestamp = dt.fromtimestamp(kt.qua_time)
+
+            q.timestamp = parse(self).timestamp(kt.qua_time)
             q.threat = kt.mal_type
             q.path = kt.filename[4:]
             q.malfile = kt.mal_file
             quarfiles[str(metafile)] = q
             kt.close()
 
         return quarfiles
-
-    def parse_from_fs(self, data=None):
-        pass

maldump/parsers/eset_parser.py

@@ -15,6 +15,7 @@
 
 from __future__ import annotations
 
+import logging
 import re
 import typing
 from pathlib import Path
@@ -24,10 +25,15 @@
 from maldump.parsers.kaitai.eset_virlog_parser import EsetVirlogParser
 from maldump.structures import Parser, QuarEntry
 from maldump.utils import DatetimeConverter as DTC
+from maldump.utils import Logger as log
+from maldump.utils import Parser as parse
+from maldump.utils import Reader as read
 
 if typing.TYPE_CHECKING:
     from datetime import datetime
 
+logger = logging.getLogger(__name__)
+
 __author__ = "Ladislav Baco"
 __copyright__ = "Copyright (C) 2017"
 __credits__ = "Ladislav Baco"
@@ -37,6 +43,7 @@
 __status__ = "Development"
 
 
+@log.log(lgr=logger)
 def parseRecord(record: dict):
     return {
         "timestamp": record.get("timestamp"),
@@ -51,6 +58,7 @@ def parseRecord(record: dict):
     }
 
 
+@log.log(lgr=logger)
 def convertToDict(parser: EsetVirlogParser):
     return [
         {
@@ -64,9 +72,20 @@ def convertToDict(parser: EsetVirlogParser):
     ]
 
 
+@log.log(lgr=logger)
 def mainParsing(virlog_path):
-    parser = EsetVirlogParser.from_file(filename=virlog_path)
-    threats = convertToDict(parser)
+    kt = parse(EsetParser).kaitai(EsetVirlogParser, virlog_path)
+    if kt is None:
+        logger.warning("Skipping virlog.dat parsing")
+        return []
+    kt.close()
+
+    threats = convertToDict(kt)
+
+    parsedRecords = []
+    for idx, record in enumerate(threats):
+        logger.debug("Parsing raw record %s/%s", idx + 1, len(threats))
+        parsedRecords.append(parseRecord(record))
 
     return [parseRecord(record) for record in threats]
 
@@ -80,37 +99,44 @@ def __init__(self):
         )
         self.regex_entry = re.compile(r"([0-9a-fA-F]+)\.NQF$")
 
+    @log.log(lgr=logger)
     def _decrypt(self, data: bytes) -> bytes:
         return bytes([((b - 84) % 256) ^ 0xA5 for b in data])
 
+    @log.log(lgr=logger)
     def _get_malfile(self, username: str, sha1: str) -> bytes:
         quarfile = self.quarpath.format(username=username)
         quarfile = Path(quarfile) / (sha1.upper() + ".NQF")
-        try:
-            with open(quarfile, "rb") as f:
-                data = f.read()
-                decrypted_data = self._decrypt(data)
-        except OSError:
-            # logging
-            print("Eset Error: could not read file", quarfile)
 
-        return decrypted_data
+        data = read.contents(quarfile, filetype="malware")
+        if data is None:
+            return b""
+
+        return self._decrypt(data)
 
+    @log.log(lgr=logger)
     def _get_metadata(self, path: Path, objhash: str) -> KaitaiParserMetadata | None:
         # metadata file has .NDF extension
         metadata_path = path / (objhash + ".NDF")
         if not metadata_path.is_file():
+            logger.debug("Metadata file not found")
+            return None
+
+        kt = parse(self).kaitai(KaitaiParserMetadata, metadata_path)
+        if kt is None:
             return None
 
-        kt = KaitaiParserMetadata.from_file(metadata_path)
         kt.close()
         return kt
 
     def parse_from_log(self, _=None) -> dict[tuple[str, datetime], QuarEntry]:
+        logger.info("Parsing from log in %s", self.name)
         quarfiles: dict[tuple[str, datetime], QuarEntry] = {}
 
-        for metadata in mainParsing(self.location):
+        for idx, metadata in enumerate(mainParsing(self.location)):
+            logger.debug("Parsing entry, idx %s", idx)
             if metadata["user"] == "SYSTEM":
+                logger.debug("Entry's (idx %s) user is SYSTEM, skipping", idx)
                 continue
             q = QuarEntry()
             q.timestamp = metadata["timestamp"]
@@ -124,26 +150,34 @@ def parse_from_log(self, _=None) -> dict[tuple[str, datetime], QuarEntry]:
     def parse_from_fs(
         self, data: dict[tuple[str, datetime], QuarEntry] | None = None
     ) -> dict[tuple[str, datetime], QuarEntry]:
+        logger.info("Parsing from filesystem in %s", self.name)
         quarfiles = {}
 
         actual_path = Path("Users/")
-        for entry in actual_path.glob(
-            "*/AppData/Local/ESET/ESET Security/Quarantine/*.NQF"
+        for idx, entry in enumerate(
+            actual_path.glob("*/AppData/Local/ESET/ESET Security/Quarantine/*.NQF")
         ):
+            logger.debug('Parsing entry, idx %s, path "%s"', idx, entry)
             res_path = re.match(self.regex_entry, entry.name)
             res_user = re.match(self.regex_user, str(entry))
 
             if not res_path:
+                logger.debug(
+                    "Entry's (idx %s) filename of incorrect format, skipping", idx
+                )
                 continue
 
             user = res_user.group(1)
             objhash = res_path.group(1)
 
             if (objhash.lower(), user) in data:
+                logger.debug("Entry (idx %s) already found, skipping", idx)
                 continue
 
-            entry_stat = entry.stat()
-
+            entry_stat = parse(self).entry_stat(entry)
+            if entry_stat is None:
+                logger.debug('Skipping entry idx %s, path "%s"', idx, entry)
+                continue
             timestamp = DTC.get_dt_from_stat(entry_stat)
             path = str(entry)
             sha1 = None

maldump/parsers/forticlient_parser.py

@@ -1,28 +1,42 @@
 from __future__ import annotations
 
+import logging
 from datetime import datetime as dt
 
 from maldump.parsers.kaitai.forticlient_parser import ForticlientParser as KaitaiParser
 from maldump.structures import Parser, QuarEntry
+from maldump.utils import Logger as log
+from maldump.utils import Parser as parse
+
+logger = logging.getLogger(__name__)
 
 
 class ForticlientParser(Parser):
+    @log.log(lgr=logger)
     def _normalize_path(self, path):
         if path[2:4] == "?\\":
             path = path[4:]
         return path
 
+    @log.log(lgr=logger)
     def _get_time(self, ts):
         return dt(ts.year, ts.month, ts.day, ts.hour, ts.minute, ts.second)
 
     def parse_from_log(self, data=None):
         pass
 
     def parse_from_fs(self, _=None) -> dict[str, QuarEntry]:
+        logger.info("Parsing from log in %s", self.name)
         quarfiles = {}
 
-        for metafile in self.location.glob("*[!.meta]"):
-            kt = KaitaiParser.from_file(metafile)
+        for idx, metafile in enumerate(self.location.glob("*[!.meta]")):
+            logger.debug('Parsing entry, idx %s, path "%s"', idx, metafile)
+
+            kt = parse(self).kaitai(KaitaiParser, metafile)
+            if kt is None:
+                logger.debug('Skipping entry idx %s, path "%s"', idx, metafile)
+                continue
+
             q = QuarEntry()
             q.timestamp = self._get_time(kt.timestamp)
             q.threat = kt.mal_type