#!/usr/bin/python3
#
# Merge a separate CVE file (such as data/next-point-update.txt) back into
# the main one.
#
# Copyright © 2020 Emilio Pozuelo Monfort <pochu@debian.org>
# Copyright (c) 2021-2022 Neil Williams <codehelp@debian.org>

import os
import sys

import setup_paths  # noqa
from debian_support import internRelease
from sectracker.parsers import (
    Bug,
    cvelist,
    writecvelist,
    PackageAnnotation,
    FlagAnnotation,
    StringAnnotation,
    XrefAnnotation
)

def merge_notes(bug, notes):
    """
    Special support for StringAnnotations.

    notes is a dict containing a list of string annotations for
    each CVE in the file being merged. Pick out the string annotations
    for this bug, ignore if already exist, append if new.
    """
    new_notes = []
    cve = bug.header.name
    merge_list = notes.get(cve)  # list of notes to merge
    if not merge_list:
        # nothing to merge
        return bug
    tagged_notes = [note.description for note in merge_list]
    bug_notes = [ann.description for ann in bug.annotations if isinstance(ann, StringAnnotation)]
    # get the list items in tagged_notes which are not in bug_notes
    new_strings = list(set(tagged_notes) - set(bug_notes))
    if not new_strings:
        return bug
    for new_ann in merge_list:
        if new_ann.description in new_strings:
            new_notes.append(new_ann)
    bug_list = list(bug.annotations)
    bug_list.extend(new_notes)
    mod_bug = Bug(
        bug.file, bug.header, tuple(bug_list)
    )
    return mod_bug


def merge_annotations(annotations, new_annotation):
    if not isinstance(new_annotation, PackageAnnotation):
        raise NotImplementedError(f"unsupported annotation of type {new_annotation.type} (line {new_annotation.line})")

    annotations = list(annotations)

    annotations_for_pkg = [ann for ann in annotations \
                           if isinstance(ann, PackageAnnotation) \
                           and ann.package == new_annotation.package]
    if not annotations_for_pkg:
        if new_annotation.release:
            raise ValueError(f"new annotation for {new_annotation.package}/{new_annotation.release} "
                              "but there is no annotation for sid")
        # new package, add it at the top
        for idx, annotation in enumerate(annotations):
            if isinstance(annotation, FlagAnnotation) \
              or isinstance(annotation, XrefAnnotation):
                continue

            annotations.insert(idx, new_annotation)
            return annotations


    # append/substitute the new one at the right place
    for idx, annotation in enumerate(annotations):
        if not isinstance(annotation, PackageAnnotation) \
          or annotation.package != new_annotation.package:
            continue

        # if the annotation is for the same package/release, replace it
        if annotation.package == new_annotation.package \
          and annotation.release == new_annotation.release:
            annotations[idx] = new_annotation
            break

        # if the next annotation's release is the same, we continue to replace
        # it in the next iteration. otherwise if we found the right place, we
        # insert the new annotation
        next_annotation = annotations[idx + 1] if len(annotations) > (idx + 1) else None
        if next_annotation and isinstance(next_annotation, PackageAnnotation) \
          and next_annotation.package == new_annotation.package \
          and internRelease(new_annotation.release) <= internRelease(next_annotation.release):
            continue

        annotations.insert(idx + 1, new_annotation)
        break

    return annotations

def parse_list(path):
    data, messages = cvelist(path)

    for m in messages:
        sys.stderr.write(str(m) + "\n")

    return data

if len(sys.argv) not in (2, 3):
    print(f"Usage: {os.path.basename(sys.argv[0])} (CVE/list) extra-cve-list")
    sys.exit(1)

if len(sys.argv) == 3:
    main_list = sys.argv[1]
else:
    main_list = os.path.dirname(__file__) + '/../data/CVE/list'

extra_list = sys.argv[-1]

data = parse_list(main_list)
extra_data = parse_list(extra_list)

for extra_bug in extra_data:
    bug = next(bug for bug in data if bug.header.name == extra_bug.header.name)

    notes = {}
    new_annotations = bug.annotations
    for extra_annotation in extra_bug.annotations:
        if isinstance(extra_annotation, FlagAnnotation):
            continue
        if isinstance(extra_annotation, StringAnnotation):
            cve = f"{extra_bug.header.name}"
            note_tag = notes.setdefault(cve, [])
            note_tag.append(extra_annotation)
            continue
        new_annotations = merge_annotations(new_annotations, extra_annotation)

    bug = bug._replace(annotations=new_annotations)
    bug = merge_notes(bug, notes)
    data = [bug if bug.header.name == old_bug.header.name else old_bug for old_bug in data]

with open(main_list, 'w') as f:
    writecvelist(data, f)

# check for and erase an .xpck file built from the merge
xpck = f"{extra_list}.xpck"
if os.path.exists(xpck):
    os.unlink(xpck)