Merge pull request #81 from ajnelson-nist/review_tree_mft_py_with_myp…

…y_strict

Makefile

@@ -59,7 +59,8 @@ check-mypy: \
 	    indxparse/INDXFind.py \
 	    indxparse/MFTINDX.py \
 	    indxparse/__init__.py \
-	    indxparse/list_mft.py
+	    indxparse/list_mft.py \
+	    indxparse/tree_mft.py
 
 check-third_party:
 	$(MAKE) \

indxparse/BinaryParser.py

@@ -24,7 +24,6 @@
 #
 #   Version v.0.1
 import array
-import mmap
 import pickle
 import struct
 import sys
@@ -35,27 +34,6 @@
 verbose = False
 
 
-class Mmap(object):
-    """
-    Convenience class for opening a read-only memory map for a file path.
-    """
-
-    def __init__(self, filename: str) -> None:
-        super(Mmap, self).__init__()
-        self._filename = filename
-        self._f = None
-        self._mmap = None
-
-    def __enter__(self):
-        self._f = open(self._filename, "rb")
-        self._mmap = mmap.mmap(self._f.fileno(), 0, access=mmap.ACCESS_READ)
-        return self._mmap
-
-    def __exit__(self, type, value, traceback):
-        self._mmap.close()
-        self._f.close()
-
-
 def debug(*message):
     global verbose
     if verbose:

indxparse/MFT.py

@@ -26,6 +26,7 @@
 #   Version v.1.1.8
 import array
 import logging
+import mmap
 import os
 import struct
 import sys
@@ -1586,7 +1587,7 @@ def get(self, k):
 class MFTEnumerator(object):
     def __init__(
         self,
-        buf: array.array,
+        buf: mmap.mmap,
         record_cache=None,
         path_cache=None,
     ) -> None:
@@ -1607,7 +1608,7 @@ def len(self) -> int:
         else:
             return floored_result + 1
 
-    def get_record_buf(self, record_num):
+    def get_record_buf(self, record_num: int) -> array.array:
         """
         @raises OverrunBufferException: if the record_num is beyond the end of the MFT
         """
@@ -1766,7 +1767,7 @@ class MFTTree(object):
 
     def __init__(
         self,
-        buf: array.array,
+        buf: mmap.mmap,
     ) -> None:
         super(MFTTree, self).__init__()
         self._buf = buf

indxparse/extract_mft_record_slack.py

@@ -20,21 +20,22 @@
 #
 #   Alex Nelson, NIST, contributed to this file.  Contributions of NIST
 #   are not subject to US Copyright.
+import mmap
 import sys
 
-from indxparse.BinaryParser import Mmap
 from indxparse.MFT import MFTEnumerator
 
 
 def main():
     filename = sys.argv[1]
 
-    with Mmap(filename) as buf:
-        enum = MFTEnumerator(buf)
-        for record in enum.enumerate_records():
-            slack = record.slack_data()
-            sys.stdout.write("\x00" * (1024 - len(slack)))
-            sys.stdout.write(slack)
+    with open(filename, "rb") as fh:
+        with mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ) as mm:
+            enum = MFTEnumerator(mm)
+            for record in enum.enumerate_records():
+                slack = record.slack_data()
+                sys.stdout.write("\x00" * (1024 - len(slack)))
+                sys.stdout.write(slack)
 
 
 if __name__ == "__main__":

indxparse/fuse-mft.py

@@ -7,13 +7,14 @@
 import calendar
 import errno
 import inspect
+import mmap
 import os
 import stat
 import sys
+from typing import Dict
 
 from fuse import FUSE, FuseOSError, Operations, fuse_get_context  # type: ignore
 
-from indxparse.BinaryParser import Mmap
 from indxparse.get_file_info import format_record
 from indxparse.MFT import Cache, MFTEnumerator, MFTRecord, MFTTree
 from indxparse.Progress import ProgressBarProgress
@@ -169,11 +170,11 @@ class MFTFuseOperations(Operations):
     MFTFuseOperations is a FUSE driver for NTFS MFT files.
     """
 
-    def __init__(self, root, mfttree, buf):
+    def __init__(self, root, mfttree, buf: mmap.mmap) -> None:
         self._root = root
         self._tree = mfttree
         self._buf = buf
-        self._opened_files = {}  # dict(int --> FH subclass)
+        self._opened_files: Dict[int, FH] = {}
 
         record_cache = Cache(1024)
         path_cache = Cache(1024)
@@ -402,11 +403,12 @@ def fsync(self, path, fdatasync, fh):
 def main():
     mft_filename = sys.argv[1]
     mountpoint = sys.argv[2]
-    with Mmap(mft_filename) as buf:
-        tree = MFTTree(buf)
-        tree.build(progress_class=ProgressBarProgress)
-        handler = MFTFuseOperations(mountpoint, tree, buf)
-        FUSE(handler, mountpoint, foreground=True)
+    with open(mft_filename, "rb") as fh:
+        with mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ) as mm:
+            tree = MFTTree(mm)
+            tree.build(progress_class=ProgressBarProgress)
+            handler = MFTFuseOperations(mountpoint, tree, mm)
+            FUSE(handler, mountpoint, foreground=True)
 
 
 if __name__ == "__main__":

indxparse/get_file_info.py

@@ -8,13 +8,13 @@
 import array
 import datetime
 import logging
+import mmap
 import re
 from string import printable
 from typing import Any, Dict
 
 from jinja2 import Template
 
-from indxparse.BinaryParser import Mmap
 from indxparse.MFT import (
     ATTR_TYPE,
     MREF,
@@ -418,27 +418,28 @@ def main():
     if results.verbose:
         logging.basicConfig(level=logging.DEBUG)
 
-    with Mmap(results.mft) as buf:
-        record_cache = Cache(results.cache_size)
-        path_cache = Cache(results.cache_size)
+    with open(results.mft, "rb") as fh:
+        with mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ) as mm:
+            record_cache = Cache(results.cache_size)
+            path_cache = Cache(results.cache_size)
 
-        enum = MFTEnumerator(buf, record_cache=record_cache, path_cache=path_cache)
+            enum = MFTEnumerator(mm, record_cache=record_cache, path_cache=path_cache)
 
-        should_use_inode = False
-        try:
-            record_num = int(results.record_or_path)
-            should_use_inode = True
-        except ValueError:
             should_use_inode = False
-
-        if should_use_inode:
-            record = enum.get_record(record_num)
-            path = results.prefix + enum.get_path(record)
-            print_indx_info(record, path)
-        else:
-            path = results.record_or_path
-            record = enum.get_record_by_path(path)
-            print_indx_info(record, results.prefix + path)
+            try:
+                record_num = int(results.record_or_path)
+                should_use_inode = True
+            except ValueError:
+                should_use_inode = False
+
+            if should_use_inode:
+                record = enum.get_record(record_num)
+                path = results.prefix + enum.get_path(record)
+                print_indx_info(record, path)
+            else:
+                path = results.record_or_path
+                record = enum.get_record_by_path(path)
+                print_indx_info(record, results.prefix + path)
 
 
 if __name__ == "__main__":

indxparse/list_mft.py

@@ -37,13 +37,13 @@
 import datetime
 import json
 import logging
+import mmap
 import sys
 import types
 from typing import Any, Dict, List, Optional, Type, Union
 
 from jinja2 import Environment, Template
 
-from indxparse.BinaryParser import Mmap
 from indxparse.get_file_info import make_model
 from indxparse.MFT import (
     ATTR_TYPE,
@@ -378,46 +378,48 @@ def main() -> None:
     else:
         progress_cls = NullProgress
 
-    with Mmap(results.filename) as buf:
-        record_cache = Cache(results.cache_size)
-        path_cache = Cache(results.cache_size)
-
-        enum = MFTEnumerator(buf, record_cache=record_cache, path_cache=path_cache)
-        progress = progress_cls(enum.len())
-        if use_default_output:
-            for record, record_path in enum.enumerate_paths():
-                output_mft_record(enum, record, results.prefix[0])
-                progress.set_current(record.inode)
-        elif results.json:
-
-            class MFTEncoder(json.JSONEncoder):
-                def default(self, obj: Any) -> Any:
-                    if isinstance(obj, datetime.datetime):
-                        return obj.isoformat("T") + "Z"
-                    elif isinstance(obj, types.GeneratorType):
-                        return [o for o in obj]
-                    return json.JSONEncoder.default(self, obj)
-
-            print("[")
-            record_count = 0
-            for record, record_path in enum.enumerate_paths():
-                if record_count > 0:
-                    print(",")
-                record_count += 1
-                m = make_model(record, record_path)
-                print(json.dumps(m, cls=MFTEncoder, indent=2))
-                progress.set_current(record.inode)
-            print("]")
-        else:
-            for record, record_path in enum.enumerate_paths():
-                sys.stdout.write(
-                    template.render(
-                        record=make_model(record, record_path), prefix=results.prefix[0]
+    with open(results.filename, "rb") as fh:
+        with mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ) as mm:
+            record_cache = Cache(results.cache_size)
+            path_cache = Cache(results.cache_size)
+
+            enum = MFTEnumerator(mm, record_cache=record_cache, path_cache=path_cache)
+            progress = progress_cls(enum.len())
+            if use_default_output:
+                for record, record_path in enum.enumerate_paths():
+                    output_mft_record(enum, record, results.prefix[0])
+                    progress.set_current(record.inode)
+            elif results.json:
+
+                class MFTEncoder(json.JSONEncoder):
+                    def default(self, obj: Any) -> Any:
+                        if isinstance(obj, datetime.datetime):
+                            return obj.isoformat("T") + "Z"
+                        elif isinstance(obj, types.GeneratorType):
+                            return [o for o in obj]
+                        return json.JSONEncoder.default(self, obj)
+
+                print("[")
+                record_count = 0
+                for record, record_path in enum.enumerate_paths():
+                    if record_count > 0:
+                        print(",")
+                    record_count += 1
+                    m = make_model(record, record_path)
+                    print(json.dumps(m, cls=MFTEncoder, indent=2))
+                    progress.set_current(record.inode)
+                print("]")
+            else:
+                for record, record_path in enum.enumerate_paths():
+                    sys.stdout.write(
+                        template.render(
+                            record=make_model(record, record_path),
+                            prefix=results.prefix[0],
+                        )
+                        + "\n"
                     )
-                    + "\n"
-                )
-                progress.set_current(record.inode)
-        progress.set_complete()
+                    progress.set_current(record.inode)
+            progress.set_complete()
 
 
 if __name__ == "__main__":

indxparse/tree_mft.py

@@ -10,27 +10,6 @@
 from indxparse.MFT import Cache, MFTTree, MFTTreeNode
 
 
-class Mmap(object):
-    """
-    Convenience class for opening a read-only memory map for a file path.
-    """
-
-    def __init__(self, filename):
-        super(Mmap, self).__init__()
-        self._filename = filename
-        self._f = None
-        self._mmap = None
-
-    def __enter__(self):
-        self._f = open(self._filename, "rb")
-        self._mmap = mmap.mmap(self._f.fileno(), 0, access=mmap.ACCESS_READ)
-        return self._mmap
-
-    def __exit__(self, type, value, traceback):
-        self._mmap.close()
-        self._f.close()
-
-
 def main() -> None:
     parser = argparse.ArgumentParser(description="Parse MFT " "filesystem structures.")
     parser.add_argument(
@@ -52,19 +31,20 @@ def main() -> None:
     if results.verbose:
         logging.basicConfig(level=logging.DEBUG)
 
-    with Mmap(results.filename) as buf:
-        record_cache = Cache(results.cache_size)
-        path_cache = Cache(results.cache_size)
+    with open(results.filename, "rb") as fh:
+        with mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ) as mm:
+            record_cache = Cache(results.cache_size)
+            path_cache = Cache(results.cache_size)
 
-        tree = MFTTree(buf)
-        tree.build(record_cache=record_cache, path_cache=path_cache)
+            tree = MFTTree(mm)
+            tree.build(record_cache=record_cache, path_cache=path_cache)
 
-        def rec(node: MFTTreeNode, prefix: str) -> None:
-            print(prefix + node.get_filename())
-            for child in node.get_children_nodes():
-                rec(child, prefix + "  ")
+            def rec(node: MFTTreeNode, prefix: str) -> None:
+                print(prefix + node.get_filename())
+                for child in node.get_children_nodes():
+                    rec(child, prefix + "  ")
 
-        rec(tree.get_root(), "")
+            rec(tree.get_root(), "")
 
 
 if __name__ == "__main__":