Mika Ayenson
78 lines
2.6 KiB
Python
78 lines
2.6 KiB
Python
# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
|
|
# or more contributor license agreements. Licensed under the Elastic License
|
|
# 2.0; you may not use this file except in compliance with the Elastic License
|
|
# 2.0.
|
|
|
|
import argparse
|
|
import difflib
|
|
import importlib
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from . import get_available_tests
|
|
from .common import CURRENT_OS
|
|
|
|
|
|
DELAY = 1
|
|
RTA_PLATFORM_TYPES = ["windows", "linux", "macos"]
|
|
|
|
|
|
def run_all():
|
|
"""Run a single RTA."""
|
|
errors = []
|
|
for ttp_file in get_available_tests(os_filter=CURRENT_OS):
|
|
print(f"---- {Path(ttp_file).name} ----")
|
|
p = subprocess.Popen([sys.executable, "-m", "rta", "-n", ttp_file])
|
|
p.wait()
|
|
code = p.returncode
|
|
|
|
if p.returncode:
|
|
errors.append((ttp_file, code))
|
|
|
|
time.sleep(DELAY)
|
|
print("")
|
|
|
|
return len(errors)
|
|
|
|
|
|
def run(ttp_name: str, *args):
|
|
"""Run all RTAs compatible with OS."""
|
|
ttp_names = sorted(get_available_tests())
|
|
if ttp_name not in ttp_names:
|
|
suggestion = ', '.join(difflib.get_close_matches(ttp_name, ttp_names, n=3))
|
|
if suggestion:
|
|
suggestion = f"Did you mean {suggestion}?"
|
|
raise ValueError(f"Unknown RTA {ttp_name}. {suggestion}")
|
|
|
|
module = importlib.import_module("rta." + ttp_name)
|
|
return module.main(*args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser("rta")
|
|
parser.add_argument("-n", "--name", dest="name", nargs='+',
|
|
help="Name(s) of test(s) to execute. E.g. bitsadmin_execution adobe_hijack")
|
|
parser.add_argument("-l", "--list", dest="list", action="store_true", help="Print a list of available tests")
|
|
parser.add_argument("-o", "--os-filter", dest="os_filter", default="all", choices=RTA_PLATFORM_TYPES,
|
|
help="Filter rule coverage summary by OS. (E.g. windows) Default: all",)
|
|
parser.add_argument("--run-all", action="store_true")
|
|
parser.add_argument("--delay", type=int, help="For run-all, the delay between executions")
|
|
parsed_args, remaining = parser.parse_known_args()
|
|
|
|
if parsed_args.name:
|
|
if parsed_args.run_all:
|
|
raise ValueError(f"Pass ttp --name or --run-all, not both")
|
|
else:
|
|
for rta_test in parsed_args.name:
|
|
rta_name = Path(rta_test).stem
|
|
exit(run(rta_name, *remaining))
|
|
|
|
elif parsed_args.list:
|
|
get_available_tests(print_list=True, os_filter=parsed_args.os_filter)
|
|
elif parsed_args.run_all:
|
|
exit(run_all())
|
|
else:
|
|
print("Execute 'python -m rta -h' to see available options")
|