#!/usr/bin/env python3

# SPDX-License-Identifier: CC0-1.0
#
# Copyright (C) 2024 Wojtek Kosior <koszko@koszko.org>

import yaml
import sys

def read_APT_data(yaml_path):
    if yaml_path:
        with open(yaml_path) as inp:
            return yaml.safe_load(inp)
    else:
        return yaml.safe_load(sys.stdin)

profiles_data = read_APT_data(None if len(sys.argv) < 2 else sys.argv[1])

groups_by_origin = {}
groups_by_techniques_by_origin = {}

for group in profiles_data["groups"]:
    origin = group["origin"]
    groups_by_origin[origin] = groups_by_origin.get(origin, []) + [group]

    groups_by_technique = groups_by_techniques_by_origin.get(origin, {})
    for tid in group["technique_ids"]:
        groups_by_technique[tid] = groups_by_technique.get(tid, []) + [group]
    groups_by_techniques_by_origin[origin] = groups_by_technique

def group_name_list(origin, tid):
    return ", ".join(sorted(
        g["name"]
        for g in groups_by_techniques_by_origin[origin].get(tid, [])
    ))

def technique_use_percent(tid, origin):
    return (100 * len(groups_by_techniques_by_origin[origin].get(tid, [])) /
            len(groups_by_origin[origin]))

def technique_popularity_score(tid):
    return sum(technique_use_percent(tid, origin)
               for origin in groups_by_origin)

all_tids = sorted(set().union(*groups_by_techniques_by_origin.values()),
                  key=technique_popularity_score,
                  reverse=True)

all_origins = sorted(groups_by_origin)

technique_names = dict((t["mitre_id"], t["name"])
                       for t in profiles_data["techniques"])

print("{")
print("\\footnotesize")
print("\\rowcolors{3}{gray!20}{white!100}")
print("\\renewcommand*{\\arraystretch}{1.3}")
print("""\
\\begin{longtable}{\
p{0.3\\textwidth} \
p{0.65in} p{0.65in} p{0.65in} p{0.65in} \
>{\\centering\\arraybackslash} p{0.85in} \
}""")
print(f"technique & {' & '.join(all_origins)} & total APT count \\\\")
print("\\hline\\hline \\endhead")

for tid in all_tids:
    name = technique_names[tid]
    group_count = sum([len(groups_by_techniques_by_origin[origin].get(tid, []))
                       for origin in all_origins])

    name_href_markup = \
        f"\\href{{https://attack.mitre.org/techniques/{tid}/}}{{{name}}}"
    group_precents_markup = ' & '.join(
        f"{round(technique_use_percent(tid, origin))}\\% groups"
        for origin in all_origins
    )

    print(f"{name_href_markup} & {group_precents_markup} & {group_count} \\\\")

print("\\end{longtable}")
print("}")