Verified Commit 3cdf0b44 authored by Sebastian Endres's avatar Sebastian Endres
Browse files

Use YaspinWrapper in inject_secrets

parent 0ece43e4
......@@ -7,9 +7,10 @@ import sys
from pathlib import Path
from typing import Optional, Union
from termcolor import cprint
from termcolor import colored, cprint
from result_parser import Result
from utils import YaspinWrapper
def parse_args():
......@@ -20,18 +21,25 @@ def parse_args():
type=Path,
help="quic-interop-runner log dirs or result.json files",
)
parser.add_argument(
"--debug",
action="store_true",
help="Debug",
)
return parser.parse_args()
class SecretsInjector:
def __init__(self, specs: list[Union[Path, Result]]):
def __init__(self, specs: list[Union[Path, Result]], debug=False):
self.specs = specs
self.num_injected = 0
self.num_already_injected = 0
self.num_no_secret_found = 0
self.num_failed = 0
self._max_log_str_len = 0
self.debug = debug
self._spinner: Optional[YaspinWrapper] = None
def find_keylog_file(self, test_case_path: Path) -> Optional[Path]:
"""Find the keylog file."""
......@@ -65,14 +73,17 @@ class SecretsInjector:
if pcap_ng_path.is_file() and pcap_ng_path.stat().st_size > 0:
self.log(
f"⚒ {pcap_ng_path} already exists. Skipping.",
color="cyan",
color="grey",
)
self.num_already_injected += 1
return
try:
self.log(f"⚒ Injecting {keylog_file} into {pcap_path} -> {pcap_ng_path}.")
self.log(
f"⚒ Injecting {keylog_file} into {pcap_path} -> {pcap_ng_path}.",
color="cyan",
)
subprocess.check_call(
[
"editcap",
......@@ -158,51 +169,60 @@ class SecretsInjector:
self.inject_in_test_repetition_run_dir(repetition_path)
def log(self, *args, **kwargs):
msg: Optional[str] = None
if args or kwargs:
msg = " ".join(args).ljust(self._max_log_str_len)
cprint(msg, **kwargs, end="\n")
log_str = (
f"⚒ injected: {self.num_injected}, "
f"failed: {self.num_failed}, "
f"already injected: {self.num_already_injected}, "
f"no keylog file found: {self.num_no_secret_found}"
log_str = ", ".join(
(
colored(f"injected: {self.num_injected}, ", color="green"),
colored(f"failed: {self.num_failed}, ", color="red"),
colored(
f"already injected: {self.num_already_injected}, ", color="yellow"
),
colored(
f"no keylog file found: {self.num_no_secret_found}", color="red"
),
)
)
self._max_log_str_len = max(self._max_log_str_len, len(log_str))
cprint(log_str, attrs=["bold"], end="\r", flush=True)
if self._spinner:
if msg:
self._spinner.write(colored(msg, **kwargs))
self._spinner.text = log_str
else:
if msg:
cprint(msg, **kwargs, end="\n")
cprint(f"⚒ {log_str}", attrs=["bold"], end="\r", flush=True)
def run(self):
for spec in self.specs:
if isinstance(spec, Result):
...
self.inject_in_result(spec)
else:
self.inject_in_log_dir(spec)
with YaspinWrapper(
debug=self.debug, text="Injecting", color="green"
) as spinner:
self._spinner = spinner
for spec in self.specs:
if isinstance(spec, Result):
self.inject_in_result(spec)
else:
self.inject_in_log_dir(spec)
self.log("⚒ Done", color="green", attrs=["bold"])
spinner.ok("✔")
def main():
args = parse_args()
cli = SecretsInjector(
specs=[Result(spec) if spec.is_file() else spec for spec in args.spec],
debug=args.debug,
)
cli.run()
class HideCursor:
def __enter__(self, *args, **kwargs):
"""hide cursor"""
print("\x1b[?25l")
def __exit__(self, *args, **kwargs):
"""show cursor"""
print("\x1b[?25h")
if __name__ == "__main__":
try:
with HideCursor():
main()
main()
except KeyboardInterrupt:
sys.exit("\nQuit")
......@@ -202,6 +202,16 @@ class YaspinWrapper:
self.yaspin.write(text)
class HideCursor:
def __enter__(self, *args, **kwargs):
"""hide cursor"""
print("\x1b[?25l")
def __exit__(self, *args, **kwargs):
"""show cursor"""
print("\x1b[?25h")
def create_relpath(path1: Path, path2: Optional[Path] = None) -> Path:
"""Create a relative path for path1 relative to path2. TODO this is broken."""
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment