Tidy up coverage classes and methods

Author: elicn

qiling/extensions/coverage/formats/base.py

@@ -4,8 +4,11 @@
 #
 
 from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
 
-from qiling import Qiling
+
+if TYPE_CHECKING:
+    from qiling import Qiling
 
 
 class QlBaseCoverage(ABC):
@@ -15,25 +18,21 @@ class QlBaseCoverage(ABC):
     all the methods marked with the @abstractmethod decorator.
     """
 
+    FORMAT_NAME: str
+
     def __init__(self, ql: Qiling):
         super().__init__()
 
         self.ql = ql
 
-    @property
-    @staticmethod
-    @abstractmethod
-    def FORMAT_NAME() -> str:
-        raise NotImplementedError
-
     @abstractmethod
-    def activate(self):
+    def activate(self) -> None:
         pass
 
     @abstractmethod
-    def deactivate(self):
+    def deactivate(self) -> None:
         pass
 
     @abstractmethod
-    def dump_coverage(self, coverage_file: str):
+    def dump_coverage(self, coverage_file: str) -> None:
         pass

qiling/extensions/coverage/formats/drcov.py

@@ -3,12 +3,18 @@
 # Cross Platform and Multi Architecture Advanced Binary Emulation Framework
 #
 
-from ctypes import Structure
-from ctypes import c_uint32, c_uint16
+from __future__ import annotations
+
+from ctypes import Structure, c_uint32, c_uint16
+from typing import TYPE_CHECKING, BinaryIO
 
 from .base import QlBaseCoverage
 
 
+if TYPE_CHECKING:
+    from qiling import Qiling
+
+
 # Adapted from https://www.ayrx.me/drcov-file-format
 class bb_entry(Structure):
     _fields_ = [
@@ -29,36 +35,35 @@ class QlDrCoverage(QlBaseCoverage):
 
     FORMAT_NAME = "drcov"
 
-    def __init__(self, ql):
+    def __init__(self, ql: Qiling):
         super().__init__(ql)
 
         self.drcov_version = 2
         self.drcov_flavor = 'drcov'
         self.basic_blocks = []
         self.bb_callback = None
 
-    @staticmethod
-    def block_callback(ql, address, size, self):
-        for mod_id, mod in enumerate(ql.loader.images):
-            if mod.base <= address <= mod.end:
-                ent = bb_entry(address - mod.base, size, mod_id)
-                self.basic_blocks.append(ent)
-                break
+    def activate(self) -> None:
+        self.bb_callback = self.ql.hook_block(self.block_callback)
 
-    def activate(self):
-        self.bb_callback = self.ql.hook_block(self.block_callback, user_data=self)
+    def deactivate(self) -> None:
+        if self.bb_callback:
+            self.ql.hook_del(self.bb_callback)
 
-    def deactivate(self):
-        self.ql.hook_del(self.bb_callback)
+    def dump_coverage(self, coverage_file: str) -> None:
+        def __write_line(bio: BinaryIO, line: str) -> None:
+            bio.write(f'{line}\n'.encode())
 
-    def dump_coverage(self, coverage_file):
         with open(coverage_file, "wb") as cov:
-            cov.write(f"DRCOV VERSION: {self.drcov_version}\n".encode())
-            cov.write(f"DRCOV FLAVOR: {self.drcov_flavor}\n".encode())
-            cov.write(f"Module Table: version {self.drcov_version}, count {len(self.ql.loader.images)}\n".encode())
-            cov.write("Columns: id, base, end, entry, checksum, timestamp, path\n".encode())
+            __write_line(cov, f"DRCOV VERSION: {self.drcov_version}")
+            __write_line(cov, f"DRCOV FLAVOR: {self.drcov_flavor}")
+            __write_line(cov, f"Module Table: version {self.drcov_version}, count {len(self.ql.loader.images)}")
+            __write_line(cov, "Columns: id, base, end, entry, checksum, timestamp, path")
+
             for mod_id, mod in enumerate(self. ql.loader.images):
-                cov.write(f"{mod_id}, {mod.base}, {mod.end}, 0, 0, 0, {mod.path}\n".encode())
-            cov.write(f"BB Table: {len(self.basic_blocks)} bbs\n".encode())
-            for bb in self.basic_blocks:
+                __write_line(cov, f"{mod_id}, {mod.base}, {mod.end}, 0, 0, 0, {mod.path}")
+
+            __write_line(cov, f"BB Table: {len(self.basic_blocks)} bbs")
+
+            for bb in self.basic_blocks.values():
                 cov.write(bytes(bb))

qiling/extensions/coverage/formats/drcov_exact.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-# 
+#
 # Cross Platform and Multi Architecture Advanced Binary Emulation Framework
 #
 
@@ -17,10 +17,7 @@ class QlDrCoverageExact(QlDrCoverage):
 
     FORMAT_NAME = "drcov_exact"
 
-    def __init__(self, ql):
-        super().__init__(ql)
-
-    def activate(self):
+    def activate(self) -> None:
         # We treat every instruction as a block on its own.
         self.bb_callback = self.ql.hook_code(self.block_callback, user_data=self)
 

qiling/extensions/coverage/formats/ezcov.py

@@ -1,14 +1,20 @@
 #!/usr/bin/env python3
-# 
+#
 # Cross Platform and Multi Architecture Advanced Binary Emulation Framework
 #
 
+from __future__ import annotations
 from collections import namedtuple
 from os.path import basename
+from typing import TYPE_CHECKING, List
 
 from .base import QlBaseCoverage
 
 
+if TYPE_CHECKING:
+    from qiling import Qiling
+
+
 # Adapted from https://github.com/nccgroup/Cartographer/blob/main/EZCOV.md#coverage-data
 class bb_entry(namedtuple('bb_entry', 'offset size mod_id')):
     def csvline(self):
@@ -27,27 +33,24 @@ class QlEzCoverage(QlBaseCoverage):
 
     FORMAT_NAME = "ezcov"
 
-    def __init__(self, ql):
+    def __init__(self, ql: Qiling):
         super().__init__(ql)
+
         self.ezcov_version = 1
-        self.ezcov_flavor  = 'ezcov'
-        self.basic_blocks  = []
-        self.bb_callback   = None
+        self.ezcov_flavor = 'ezcov'
+        self.basic_blocks: List[bb_entry]  = []
+        self.bb_callback = None
 
-    @staticmethod
-    def block_callback(ql, address, size, self):
-        mod = ql.loader.find_containing_image(address)
-        if mod is not None:
-            ent = bb_entry(address - mod.base, size, basename(mod.path))
-            self.basic_blocks.append(ent)
+    def block_callback(self, ql: Qiling, address: int, size: int):
 
-    def activate(self):
-        self.bb_callback = self.ql.hook_block(self.block_callback, user_data=self)
+    def activate(self) -> None:
+        self.bb_callback = self.ql.hook_block(self.block_callback)
 
-    def deactivate(self):
-        self.ql.hook_del(self.bb_callback)
+    def deactivate(self) -> None:
+        if self.bb_callback:
+            self.ql.hook_del(self.bb_callback)
 
-    def dump_coverage(self, coverage_file):
+    def dump_coverage(self, coverage_file: str) -> None:
         with open(coverage_file, "w") as cov:
             cov.write(f"EZCOV VERSION: {self.ezcov_version}\n")
             cov.write("# Qiling EZCOV exporter tool\n")

qiling/extensions/tracing/formats/base.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-# 
+#
 # Cross Platform and Multi Architecture Advanced Binary Emulation Framework
 # This code structure is copied and modified from the coverage extension
 
@@ -12,24 +12,20 @@ class QlBaseTrace(ABC):
     To add support for a new coverage format, just derive from this class and implement
     all the methods marked with the @abstractmethod decorator.
     """
-    
+
+    FORMAT_NAME: str
+
     def __init__(self):
         super().__init__()
 
-    @property
-    @staticmethod
-    @abstractmethod
-    def FORMAT_NAME():
-        raise NotImplementedError
-
     @abstractmethod
-    def activate(self):
+    def activate(self) -> None:
         pass
 
     @abstractmethod
-    def deactivate(self):
+    def deactivate(self) -> None:
         pass
 
     @abstractmethod
-    def dump_trace(self, trace_file):
+    def dump_trace(self, trace_file: str) -> None:
         pass