#!/usr/bin/python3
#
# Update xrefs in data/CVE/list
#
# Copyright © 2023 Emilio Pozuelo Monfort <pochu@debian.org>

import argparse
import os

import setup_paths  # noqa
from sectracker import parsers

def get_annotation(annotations, ann_type):
    for ann in annotations:
        if isinstance(ann, ann_type):
            return ann


def add_xref(cve, xref):
    ann = get_annotation(cve.annotations, parsers.XrefAnnotation)
    if not ann:
        ann = parsers.XrefAnnotation(0, "xref", list())
        # TODO: annotations order is important atm
        flag_ann = get_annotation(cve.annotations, parsers.FlagAnnotation)
        idx = cve.annotations.index(flag_ann) + 1 if flag_ann else 0
        cve.annotations.insert(idx, ann)

    if not xref in ann.bugs:
        ann.bugs.append(xref)


def parse_advfile(filename, parsemethod):
    global cve_map

    advs = parsemethod(filename)

    for adv in advs:
        for ann in adv.annotations:
            if isinstance(ann, parsers.XrefAnnotation):
                cves = ann.bugs
                for cvename in cves:
                    if not cvename.startswith('CVE-'):
                        continue

                    cve = cve_map[cvename]
                    add_xref(cve, adv.header.name)


def parse_dsafile(dsafile):
    return parse_advfile(dsafile, parsers.dsalist)


def parse_dtsafile(dtsafile):
    return parse_advfile(dtsafile, parsers.dtsalist)


def parse_dlafile(dlafile):
    return parse_advfile(dlafile, parsers.dlalist)


def remove_xrefs(cves):
    for cve in cves:
        #cve.annotations = [ann
        #                   for ann in cve.annotations
        #                   if not isinstance(ann, parsers.XrefAnnotation)]

        ann = get_annotation(cve.annotations, parsers.XrefAnnotation)

        if ann:
            # we have CVE- cross-references, keep those and remove
            # the rest, which we will re-add later if appropriate
            ann.bugs = [bug for bug in ann.bugs if bug.startswith('CVE-')]
            if len(ann.bugs) == 0:
                cve.annotations.remove(ann)


default_workdir = os.path.join(os.path.dirname(os.path.dirname(__file__)))

parser = argparse.ArgumentParser(description='Update cross-references in CVE list')
parser.add_argument('--work-dir', help='path to security-tracker repo (default: relative to the script)', default=default_workdir)
args = parser.parse_args()

dsa_list = args.work_dir + '/data/DSA/list'
dtsa_list = args.work_dir + '/data/DTSA/list'
dla_list = args.work_dir + '/data/DLA/list'
main_list = args.work_dir + '/data/CVE/list'

cves = parsers.cvelist(main_list)
cve_map = {cve.header.name: cve for cve in cves}

# We remove the Xrefs, then re-parse the various advisory files and re-add
# them
remove_xrefs(cves)
parse_dsafile(dsa_list)
parse_dtsafile(dtsa_list)
parse_dlafile(dla_list)

# write the CVE file back
with open(main_list, 'w') as f:
    parsers.writecvelist(cves, f)