main.py

import argparse
from argparse import Namespace
from pipeline import pipeline_init_cm, pipeline_injection, pipeline_detection, pipeline_assessor, pipeline_evaluation


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--scenario", type = str, default = "STD")

    args, unknown = parser.parse_known_args()
    args_dict = vars(args).copy()

    # parse unknown arguments
    for i in range(0, len(unknown), 2):
        if unknown[i].startswith("--") and (i + 1) < len(unknown):
            key = unknown[i][2:]
            value = unknown[i + 1]

            if value.isdigit():
                value = int(value)
            elif value.lower() in ["true", "false"]:
                value = value.lower() == "true"

            args_dict[key] = value

    return Namespace(**args_dict)


def main(args):
    cm = pipeline_init_cm(args = args)  # init config manager

    # Injection Phase: Generate poisoned files based on scenario configuration.
    pipeline_injection(cm = cm, flag = cm.global_config.Flag_injection)

    # Detection Phase: Run selected detection models and obtain classification results (detection side, micro-level).
    pipeline_detection(cm = cm, flag = cm.global_config.Flag_detection)

    # Assessment Phase: Evaluate the impact of generated files and compute distributional distances (attack side, micro-level).
    pipeline_assessor(cm = cm, flag = cm.global_config.Flag_assess)

    # Evaluation Phase: Conduct macro-level evaluation per scenario.
    pipeline_evaluation(cm = cm, flag = cm.global_config.Flag_evaluation)


if __name__ == "__main__":
    args = parse_args()
    main(args)