Add main.py

main.py

+import sys
+import os
+import stat
+import traceback
+import hashlib
+import io
+import multiprocessing
+import json
+
+
+class Scheduler:
+    def __init__(self):
+        self.mp_pool = multiprocessing.Pool()
+
+    def submit(self, f, *args):
+        return self.mp_pool.apply_async(f, args)
+
+    def join(self):
+        self.mp_pool.close()
+        self.mp_pool.join()
+
+
+
+def stat_to_dict(stat_result):
+    return {
+        'mode': getattr(stat_result, 'st_mode', None),
+        'ino': getattr(stat_result, 'st_ino', None),
+        'dev': getattr(stat_result, 'st_dev', None),
+        'nlink': getattr(stat_result, 'st_nlink', None),
+        'uid': getattr(stat_result, 'st_uid', None),
+        'gid': getattr(stat_result, 'st_gid', None),
+        'size': getattr(stat_result, 'st_size', None),
+        'atime': getattr(stat_result, 'st_atime', None),
+        'mtime': getattr(stat_result, 'st_mtime', None),
+        'ctime': getattr(stat_result, 'st_ctime', None),
+        'blocks': getattr(stat_result, 'st_blocks', None),
+        'blksize': getattr(stat_result, 'st_blksize', None),
+        'rdev': getattr(stat_result, 'st_rdev', None),
+        'flags': getattr(stat_result, 'st_flags', None),
+        'gen': getattr(stat_result, 'st_gen', None),
+        'birthtime': getattr(stat_result, 'st_birthtime', None),
+    }
+
+
+def Node(path):
+    try:
+        stat_result = os.stat(path)
+        if stat.S_ISDIR(stat_result.st_mode):
+            return DirNode(path, stat_result)
+        elif stat.S_ISREG(stat_result.st_mode):
+            return FileNode(path, stat_result)
+        elif stat.S_ISLNK(stat_result.st_mode):
+            return SymlinkNode(path, stat_result)
+    except Exception:
+      pass
+
+
+def compute_hash(path):
+    try:
+        with open(path, mode='rb') as f:
+            sha256 = hashlib.sha256()
+            sha512 = hashlib.sha512()
+
+            while True:
+                data = f.read(io.DEFAULT_BUFFER_SIZE)
+
+                sha256.update(data)
+                sha512.update(data)
+
+                if len(data) < io.DEFAULT_BUFFER_SIZE:
+                    break
+
+            return {
+                'sha256': sha256.hexdigest(),
+                'sha512': sha512.hexdigest(),
+            }
+    except Exception:
+        return {
+           'sha256': "<error>",
+           'sha512': "<error>",
+        }
+
+
+class DirNode:
+    def __init__(self, path, stat_result):
+        self.path = path
+        self.stat_result = stat_result
+        self.children = []
+
+    def build_node(self):
+        with os.scandir(self.path) as entries:
+            for entry in entries:
+                path = os.fsdecode(entry.path)
+                node = Node(path)
+                if node:
+                  self.children.append(node)
+                  node.build_node()
+        self.data = {'stat': self.stat_result}
+
+    def gather_node(self, result):
+        result[self.path] = self.data
+        for node in self.children:
+            node.gather_node(result)
+
+
+class FileNode:
+    def __init__(self, path, stat_result):
+        self.path = path
+        self.stat_result = stat_result
+
+    def build_node(self):
+        self.task = scheduler.submit(compute_hash, self.path)
+
+    def gather_node(self, result):
+        data = self.task.get()
+        data['stat'] = stat_to_dict(self.stat_result)
+        result[self.path] = data
+
+
+class SymlinkNode:
+    def __init__(self, path, stat_result):
+        self.path = path
+        self.stat_result = stat_result
+
+    def build_node(self):
+        self.data = {
+            'stat': stat_to_dict(stat_result),
+            'target': os.readlink(entry_path),
+        }
+
+    def gather_node(self, result):
+        result[self.path] = self.data
+
+
+def main(path='.'):
+    global scheduler
+    scheduler = Scheduler()
+
+    root = Node(path)
+    root.build_node()
+
+    scheduler.join()
+
+    result = {}
+    root.gather_node(result) 
+
+    json.dump(result, sys.stdout, indent=2, separators=(',', ': '))
+
+
+if __name__ == "__main__":
+    if len(sys.argv) == 2:
+      main(sys.argv[1])
+    else:
+      main()