Sanitises a folder of RWR profiles

sanitise.py

import sys, time, hashlib, zipfile
from collections import Counter
from pathlib import Path
import xml.etree.ElementTree as XmlET


script_dir = Path(__file__).parent
salt_file = script_dir / "salt.txt"
salt_info = "A salt is required for the hash function used to protect the player SIDs.\n" \
            "This salt must remain constant once set in order for alts to be coalesced " \
            "and profiles to be comparable across dates.\n" \
            f"It will be stored in '{salt_file}' and loaded automatically on subsequent runs.\n" \
            "It should be a decent length (12 characters minimum) alphanumeric string with symbols."
outzip_path = script_dir / f"{time.strftime('%Y%m%d_%H%M%S')}.zip"

# unnecessary attributes lists
stats_attribs = ["player_kills", "teamkills", "longest_kill_streak", "times_got_healed", "soldiers_healed",
                 "distance_moved", "targets_destroyed", "vehicles_destroyed", "shots_fired", "throwables_thrown"]
person_attribs = ["max_authority_reached", "authority", "faction", "name", "version", "alive", "soldier_group_id",
                  "soldier_group_name", "block", "squad_size_setting"]


def create_salt_file():
    print(salt_info)
    while True:
        s = input("Salt: ")
        if not s:
            print("Salt must not be empty")
            continue
        elif len(s) < 12:
            print("Salt must be 12 characters or longer")
            continue
        else:
            break
    print(f"Writing salt to '{salt_file}'...")
    salt_file.write_text(s, encoding="utf-8")


if __name__ == '__main__':
    # check the number of arguments provided to the script
    if (l := len(sys.argv)) != 2:
        print(f"Error: incorrect number of parameters provided: expected 1, got {l - 1}")
        print("Usage: sanitise.py <profiles_dir/>")
        sys.exit(1)

    # construct a Path from first argument and check it exists
    target_path = Path(sys.argv[1])
    if not target_path.exists():
        print(f"Error: '{target_path.resolve()}' does not exist")
        print("Usage: sanitise.py <profiles_dir/>")
        sys.exit(2)

    # check the target_path is a dir
    if not target_path.is_dir():
        print(f"Error: target path must be a directory")
        print("Usage: sanitise.py <profiles_dir/>")
        sys.exit(3)

    # check if the salt file exists and create it if it doesn't
    if not salt_file.exists():
        print(f"Warning: '{salt_file}' not found...")
        create_salt_file()

    # load the salt from the salt file
    salt = salt_file.read_text(encoding="utf-8").strip()

    # setup the output zip
    outzip = zipfile.ZipFile(outzip_path, mode="x", compression=zipfile.ZIP_DEFLATED, compresslevel=9)

    print(f"Processing profiles in '{target_path}'...")
    profile_paths, count, errors = target_path.glob("*.profile"), 0, 0
    t0 = time.time()
    for profile_path in profile_paths:
        try:
            id_ = profile_path.stem
            # print(f"Processing '{id_}'...")
            person_path = target_path / f"{id_}.person"
            # if associated <id_>.person doesn't exist, skip iteration
            if not person_path.exists():
                errors += 1
                print(f"Warning: '{person_path}' does not exist, skipping...")
                continue
            # load the XML elements from <id_> profile and person
            with profile_path.open(encoding="utf-8") as profile_file, person_path.open(encoding="utf-8") as person_file:
                profile_xml = XmlET.fromstring(profile_file.read())
                person_xml = XmlET.fromstring(person_file.read())
            # sanitise the profile element
            # pop digest, rid, old_digest out of the attrib dict
            # this ensures they will not be present in the output
            profile_xml.attrib.pop("digest", None)
            profile_xml.attrib.pop("rid", None)
            profile_xml.attrib.pop("old_digest", None)
            # replace the sid with hash(sid)
            sid = profile_xml.attrib["sid"]
            hash_ = hashlib.md5(f"{salt}{sid}".encode("utf-8")).hexdigest()
            profile_xml.attrib["sid"] = hash_

            # clean up the profile element - removing superfluous data
            profile_xml.attrib.pop("color", None)
            stats_elem = profile_xml.find("stats")
            if stats_elem is not None:
                # remove unnecessary attributes from the profile/stats element
                for stats_attrib in stats_attribs:
                    stats_elem.attrib.pop(stats_attrib, None)
                # remove all monitor element children
                # findall collects all matching elements before iteration begins
                for monitor_elem in stats_elem.findall("monitor"):
                    stats_elem.remove(monitor_elem)

            # clean up the person element
            # remove unnecessary attributes
            for person_attrib in person_attribs:
                person_xml.attrib.pop(person_attrib, None)
            # remove person/order element
            order_elem = person_xml.find("order")
            if order_elem is not None:
                person_xml.remove(order_elem)
            # remove person/(equipped)item elements
            for equipped_item_elem in person_xml.findall("item"):
                person_xml.remove(equipped_item_elem)

            # condense stash info
            stash_elem = person_xml.find("stash")
            if stash_elem is not None:
                stash_item_elems = stash_elem.findall("item")
                stash_items = [i.get("key") for i in stash_item_elems]
                stash_item_counter = Counter(stash_items)
                # remove the old item elements
                for stash_item_elem in stash_item_elems:
                    stash_elem.remove(stash_item_elem)
                # add the condensed data to the stash element
                for key, count_item in stash_item_counter.items():
                    i_elem = XmlET.Element("i")
                    i_elem.attrib["k"] = key
                    i_elem.attrib["c"] = str(count_item)
                    stash_elem.append(i_elem)

            # condense backpack info
            backpack_elem = person_xml.find("backpack")
            if backpack_elem is not None:
                backpack_item_elems = backpack_elem.findall("item")
                backpack_items = [i.get("key") for i in backpack_item_elems]
                backpack_item_counter = Counter(backpack_items)
                # remove the old item elements
                for backpack_item_elem in backpack_item_elems:
                    backpack_elem.remove(backpack_item_elem)
                # add the condensed data to the backpack element
                for key, count_item in backpack_item_counter.items():
                    i_elem = XmlET.Element("i")
                    i_elem.attrib["k"] = key
                    i_elem.attrib["c"] = str(count_item)
                    backpack_elem.append(i_elem)

            # make a data element so we can output the profile and corresponding person together
            data_element = XmlET.Element("data")
            data_element.append(profile_xml)
            data_element.append(person_xml)
            # write the data xml into the outzip
            data_str = XmlET.tostring(data_element, encoding="unicode")
            outzip.writestr(f"{id_}.xml", data_str, compress_type=zipfile.ZIP_DEFLATED, compresslevel=9)
            count += 1
        except Exception as e:
            errors += 1
            print(f"Warning: '{profile_path}' raised '{type(e).__name__}: {e}', skipping...")
            continue

    outzip.close()
    print(f"Processed {count} successfully, {errors} skipped due to errors, in {(time.time() - t0):.2f}s")