STIG_Manipulator.py

#!/usr/bin/env python3

import sys
from lxml import etree
from pathlib import Path
import random

stig_folder = "./tests/test_data"


def rand():
    statuses = ["Open", "NotAFinding", "NotApplicable"]
    random.shuffle(statuses)
    return statuses.pop()


# accessing and printing files in directory and subdirectory
for filename in Path(stig_folder).glob(
    "**/*.ckl",
):
    tree = etree.parse(filename)
    for s in tree.xpath("//CHECKLIST/STIGS/iSTIG/VULN/STATUS"):
        if len(sys.argv) > 1 and sys.argv[1] == "shuffle":
            s.text = rand()
        elif len(sys.argv) > 1 and sys.argv[1] == "fail":
            s.text = "Open"
        elif len(sys.argv) > 1 and sys.argv[1] == "pass":
            s.text = "NotAFinding"
        else:
            sys.exit("Please provide an argument: shuffle, fail, or pass")
    with open(filename, "wb") as f:
        tree.write(f)