#!/usr/bin/env python3
# Convert a CVE file into a CVEProject/cvelist style json vulnogram
#
# Copyright 2020, Canonical Ltd.
# Author: Alex Murray <alex.murray@canonical.com>

import cve_lib
import json
import os
import sys
import time

if len(sys.argv) != 2:
    print(f"Usage: {sys.argv[0]} [CVE]", file=sys.stderr)
    sys.exit(1)

cveid = sys.argv[1]
cve = cve_lib.load_cve(cveid)
template_fn = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                           "vulnogram.json")
with open(template_fn, 'r') as template:
    js = json.load(template)

js["CVE_data_meta"]["ASSIGNER"] = "security@ubuntu.com"
# TODO: this might not be correct if this is a Linux kernel CVE or similar
js["affects"]["vendor"]["vendor_data"][0]["vendor_name"] = "Canonical"

js["CVE_data_meta"]["ID"] = cve["Candidate"]
# Add description
js["description"]["description_data"][0]["value"] = cve["Description"].strip()

public = None
try:
    public = time.strptime(cve['PublicDate'], "%Y-%m-%d %H:%M:%S %Z")
except ValueError:
  try:
    public = time.strptime(cve['PublicDate'], "%Y-%m-%d %H:%M:%S")
  except ValueError:
      try:
        public = time.strptime(cve['PublicDate'], "%Y-%m-%d")
      except ValueError:
        pass

if public is None:
  js["CVE_data_meta"]["DATE_PUBLIC"] = ""
else:
  js["CVE_data_meta"]["DATE_PUBLIC"] = time.strftime("%Y-%m-%dT%H:%M:%S.000Z",
                                                   public)
js["CVE_data_meta"]["TITLE"] = "TODO - INSERT TITLE HERE"

# if a USN is listed in references, use this as the CNA advisory URL
for ref in cve["References"].splitlines():
    if "usn.ubuntu.com" in ref:
        js["source"]["advisory"] = ref
    # add all references
    if ref.strip() != "":
        js["references"]["reference_data"].append({"refsource": "CONFIRM", "url": ref})
# if any launchpad bugs listed then use as defect
js["source"]["defect"] = []
for bug in cve["Bugs"].splitlines():
    if "bugs.launchpad.net" in bug:
        js["source"]["defect"].append(bug)
    # add all bugs as references
    if bug.strip() != "":
        js["references"]["reference_data"].append({"refsource": "CONFIRM", "url": bug})

# remove any blank references from the template
for ref in js["references"]["reference_data"]:
     if "url" not in ref or ref["url"] == "":
          js["references"]["reference_data"].remove(ref)


# add the affected product(s) and versions
for pkg in cve["pkgs"].keys():
    product = {"product_name": pkg,
               "version": { "version_data": []}}
    desc = "This issue affects:\n" + pkg
    for release in cve["pkgs"][pkg].keys():
        if cve["pkgs"][pkg][release][0] == 'released':
            version_name = cve["pkgs"][pkg][release][1].split("-")[0]
            version_value = cve["pkgs"][pkg][release][1]
            product["version"]["version_data"].append({"version_affected": "<",
                                                       "version_name": version_name,
                                                       "version_value": version_value})
            # also add text describing affected version to description as well since
            # MITRE likes this redundancy
            desc += "\n" + version_name + " versions prior to " + version_value + ";"
    js["description"]["description_data"].append({"lang": "eng", "value": desc})

    js["affects"]["vendor"]["vendor_data"][0]["product"]["product_data"].append(product)

# remove any empty products
for product in js["affects"]["vendor"]["vendor_data"][0]["product"]["product_data"]:
    if "product_name" not in product or product["product_name"] == "":
        js["affects"]["vendor"]["vendor_data"][0]["product"]["product_data"].remove(product)


# Add a credit if we have one
if cve["Discovered-by"].strip() != "":
    js["credit"] = [{"lang": "eng", "value": cve["Discovered-by"].strip()}]

# Add a workaround if we have one
try:
    if cve["Mitigation"].strip() != "":
        js["work_around"] = [{"lang": "eng", "value": cve["Mitigation"].strip()}]
except:
    pass

# Add CVSS if we have defined one
if isinstance(cve["CVSS"], list):
    for (cvss_source, cvss_entry) in cve["CVSS"]:
        if cvss_source.strip() == 'ubuntu':
            js["impact"]["cvss"] = cve_lib.parse_cvss(cvss_entry)["baseMetricV3"]["cvssV3"]
elif cve["CVSS"].strip() != "":
    js["impact"]["cvss"] = cve_lib.parse_cvss(cve["CVSS"])["baseMetricV3"]["cvssV3"]

# Finally fill in a todo for the problem type
js["problemtype"]["problemtype_data"][0]["description"][0]["value"] = "TODO"

print(json.dumps(js, indent=2))