sigma-rules/rta/__main__.py

# 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")