Add noise argument for obfuscation (#278)

* add noise for obfuscation

* fix

Author: babenek

download_data.py

@@ -239,12 +239,12 @@ def check_snapshot_meta(snapshot: dict) -> int:
     return result
 
 
-def main(args: Namespace):
-    if os.path.exists(args.data_dir):
-        if not args.clean_data:
-            raise FileExistsError(f"{args.data_dir} directory already exists. "
+def process(arguments: Namespace) -> int:
+    if os.path.exists(arguments.data_dir):
+        if not arguments.clean_data:
+            raise FileExistsError(f"{arguments.data_dir} directory already exists. "
                                   f"Please remove it or select other directory.")
-        shutil.rmtree(args.data_dir)
+        shutil.rmtree(arguments.data_dir)
 
     with open("snapshot.json", encoding="utf_8") as f:
         snapshot = json.load(f)
@@ -253,31 +253,34 @@ def main(args: Namespace):
         logger.critical(f"Check logs, fix and restart if necessary: {new_meta_count}")
         return 1
 
-    jobs = 1 if not args.jobs else max(1, int(args.jobs))
-    if not args.skip_download:
+    jobs = 1 if not arguments.jobs else max(1, int(arguments.jobs))
+    if not arguments.skip_download:
         logger.info("Start download")
         os.makedirs(TMP_DIR, exist_ok=True)
         download(snapshot, jobs)
         logger.info("Download finished. Now processing the files...")
     else:
         logger.info("Download skipped. Now processing the files...")
 
-    removed_meta = move_files(snapshot, args.data_dir)
+    removed_meta = move_files(snapshot, arguments.data_dir)
     # check whether there were issues with downloading
     assert 0 == len(removed_meta), removed_meta
     logger.info("Finalizing dataset. Please wait a moment...")
-    obfuscate_creds("meta", args.data_dir)
-    logger.info(f"Done! All files saved to {args.data_dir}")
+    obfuscate_creds("meta", arguments.data_dir, arguments.noise)
+    logger.info(f"Done! All files saved to {arguments.data_dir}")
     return 0
 
 
-if __name__ == "__main__":
+def main(argv) -> int:
     parser = ArgumentParser(prog="python download_data.py")
-
     parser.add_argument("--data_dir", dest="data_dir", default="data", help="Dataset location after download")
     parser.add_argument("--jobs", dest="jobs", help="Jobs for multiprocessing")
     parser.add_argument("--skip_download", help="Skip download", action="store_const", const=True)
     parser.add_argument("--clean_data", help="Recreate data dir", action="store_const", const=True)
-    _args = parser.parse_args()
+    parser.add_argument("--noise", help="Seed perturbation", type=int, default=0)
+    arguments = parser.parse_args(argv[1:])
+    return process(arguments)
+
 
-    sys.exit(main(_args))
+if __name__ == """__main__""":
+    sys.exit(main(sys.argv))

obfuscate_creds.py

@@ -397,7 +397,7 @@ def gen_random_value(value):
     return obfuscated_value
 
 
-def replace_rows(data: List[MetaRow], lines: List[str]):
+def replace_rows(data: List[MetaRow], lines: List[str], noise: int):
     # Change data in already copied files
     for row in data:
         # PEM keys and other multiple-line credentials is processed in other function
@@ -418,7 +418,7 @@ def replace_rows(data: List[MetaRow], lines: List[str]):
         old_line = lines[row.LineStart - 1]
         value = old_line[row.ValueStart:row.ValueEnd]
         # CredSweeper may scan huge lines since v1.6
-        random.seed((row.ValueStart | (row.LineStart << 16)) ^ int(row.FileID, 16))
+        random.seed((row.ValueStart | (row.LineStart << 16)) ^ int(row.FileID, 16) ^ noise)
         obfuscated_value = get_obfuscated_value(value, row)
         new_line = old_line[:row.ValueStart] + obfuscated_value + old_line[row.ValueEnd:]
 
@@ -532,7 +532,7 @@ def create_new_multiline(lines: List[str], starting_position: int):
     return new_lines
 
 
-def process_pem_key(row: MetaRow, lines:List[str]):
+def process_pem_key(row: MetaRow, lines: List[str], noise: int):
     # Change data in already copied files (only keys)
     try:
         # Skip credentials that are not PEM or multiline
@@ -543,7 +543,7 @@ def process_pem_key(row: MetaRow, lines:List[str]):
             # skip double obfuscation for the categories
             return
 
-        random.seed(row.LineStart ^ int(row.FileID, 16))
+        random.seed(row.LineStart ^ int(row.FileID, 16) ^ noise)
 
         if '' != row.CryptographyKey:
             new_lines = create_new_key(lines[row.LineStart - 1:row.LineEnd])
@@ -557,13 +557,14 @@ def process_pem_key(row: MetaRow, lines:List[str]):
         logger.critical(exc)
         raise
 
-def process_pem_keys(data: List[MetaRow], lines:List[str]):
+
+def process_pem_keys(data: List[MetaRow], lines: List[str], noise: int):
     for row in data:
         if 'T' == row.GroundTruth and "Private Key" == row.Category:
-            process_pem_key(row, lines)
+            process_pem_key(row, lines, noise)
 
 
-def obfuscate_creds(meta_dir: str, dataset_dir: str):
+def obfuscate_creds(meta_dir: str, dataset_dir: str, noise: int = 0):
     dataset_files = {}
     for meta_row in read_meta(meta_dir):
         meta_row.FilePath = meta_row.FilePath.replace("data", dataset_dir, 1)
@@ -581,25 +582,28 @@ def obfuscate_creds(meta_dir: str, dataset_dir: str):
             logger.critical(exc)
             raise
         meta_rows.sort(key=lambda x: (x.LineStart, x.LineEnd, x.ValueStart, x.ValueEnd))
-        replace_rows(meta_rows, lines)
-        process_pem_keys(meta_rows, lines)
+        replace_rows(meta_rows, lines, noise)
+        process_pem_keys(meta_rows, lines, noise)
 
         with open(dataset_file, "w", encoding="utf8") as f:
             f.write('\n'.join(lines))
 
 
-def main(args: Namespace):
-    obfuscate_creds(args.meta_dir, args.data_dir)
+def obfuscate(arguments: Namespace) -> int:
+    obfuscate_creds(arguments.meta_dir, arguments.data_dir, arguments.noise)
     logger.info(f"Obfuscation was done")
     return 0
 
 
-if __name__ == "__main__":
+def main(argv) -> int:
     parser = ArgumentParser(prog="python obfuscate_creds.py")
-
     parser.add_argument("--meta_dir", dest="meta_dir", default="meta", help="Dataset markup")
     parser.add_argument("--data_dir", dest="data_dir", default="data", help="Dataset location after download")
     parser.add_argument("--jobs", dest="jobs", help="Jobs for multiprocessing")
-    _args = parser.parse_args()
+    parser.add_argument("--noise", help="Seed perturbation", type=int, default=0)
+    arguments = parser.parse_args(argv[1:])
+    return obfuscate(arguments)
+
 
-    sys.exit(main(_args))
+if __name__ == """__main__""":
+    sys.exit(main(sys.argv))