Inject all keylog files

post_process_logs.py

@@ -13,6 +13,7 @@ import os
 import shutil
 import subprocess
 import sys
+from itertools import chain
 from pathlib import Path
 from typing import Optional, Union
 
@@ -107,7 +108,7 @@ class PostProcessor:
 
             return None
 
-    def find_keylog_file(self, test_case_path: Path) -> Optional[Path]:
+    def find_keylog_files(self, test_case_path: Path) -> list[Path]:
         """Find the keylog file."""
         results = list[Path]()
 
@@ -116,16 +117,7 @@ class PostProcessor:
                 if file == "keys.log":
                     results.append(Path(root) / file)
 
-        if not results:
-            return None
-
-        if len(results) > 1:
-            self.log(
-                f"⚒ found more than one keylog file for {test_case_path}. Using the first one.",
-                color="yellow",
-            )
-
-        return results[0]
+        return results
 
     def rename_qlogs_in_test_repetition_run_dir(self, test_run_dir: Path):
         """Rename QLOG files in test repetition run dir."""
@@ -183,8 +175,8 @@ class PostProcessor:
 
             self.log()
 
-    def inject(self, pcap_path: Path, pcap_ng_path: Path, keylog_file: Path):
-        """Inject keylog file into pcap."""
+    def inject(self, pcap_path: Path, pcap_ng_path: Path, keylog_files: list[Path]):
+        """Inject keylog files into pcap."""
 
         if not pcap_path.is_file():
             self.log(
@@ -196,14 +188,18 @@ class PostProcessor:
 
         try:
             self.log(
-                f"⚒ Injecting {keylog_file} into {pcap_path} -> {pcap_ng_path}.",
+                f"⚒ Injecting {len(keylog_files)} keylog files into {pcap_path} -> {pcap_ng_path}.",
                 color="cyan",
             )
             subprocess.check_call(
                 [
                     "editcap",
-                    "--inject-secrets",
-                    f"tls,{keylog_file}",
+                    *chain(
+                        *(
+                            ["--inject-secrets", f"tls,{keylog_file}"]
+                            for keylog_file in keylog_files
+                        ),
+                    ),
                     pcap_path,
                     pcap_ng_path,
                 ]
@@ -214,7 +210,7 @@ class PostProcessor:
                 pcap_path.unlink()
         except subprocess.CalledProcessError:
             self.log(
-                f"⨯ Failed to inject secrets {keylog_file} into {pcap_path}.",
+                f"⨯ Failed to inject {len(keylog_files)} keylog files into {pcap_path}.",
                 color="red",
             )
             self.num_inject_failed += 1
@@ -225,7 +221,7 @@ class PostProcessor:
     def inject_secrets_in_test_repetition_run_dir(self, test_run_dir: Path):
         """docstring for inject_secrets_in_test_repetition_run_dir"""
 
-        keylog_file = self.find_keylog_file(test_run_dir)
+        keylog_files = self.find_keylog_files(test_run_dir)
 
         for pcap_name in ("left", "right"):
             pcap_root = test_run_dir / "sim"
@@ -245,7 +241,7 @@ class PostProcessor:
 
                 continue
 
-            if not keylog_file:
+            if not keylog_files:
                 self.log(
                     f"⨯ no keylog file found in {test_run_dir}",
                     color="red",
@@ -257,13 +253,15 @@ class PostProcessor:
             self.inject(
                 pcap_path=pcap_path,
                 pcap_ng_path=pcap_ng_path,
-                keylog_file=keylog_file,
+                keylog_files=keylog_files,
             )
 
             self.log()
 
-        if self.clean and keylog_file and keylog_file.is_file():
-            keylog_file.unlink()
+        if self.clean:
+            for keylog_file in keylog_files:
+                if keylog_file.is_file():
+                    keylog_file.unlink()
 
     def post_process_test_repetition_run_dir(self, test_run_dir: Path):
         """Inject secrets into a test repetition run log_dir."""