#!/usr/bin/python

#
# A script for looking for kernel issues that can be fixed by ReadyKernel
#
# Copyright (c) 2016-2019 Virtuozzo International GmbH
#

import sys
import os
import urllib2
import json
import platform
import re
import subprocess
import argparse
import datetime

BASEURL = "https://readykernel.com"
APIURL = BASEURL + "/api/v1"

def print_warning(e):
    print('Warning: Unable to check whether the kernel is supported ({})'.format(e.reason))

'''
Get version of the running kernel by means of "uname -r" - it is used
as unique kernel identifier in RK
'''
def get_kernel_ver():
    proc = subprocess.Popen(['uname', '-r'],  stdout=subprocess.PIPE)
    out, err = proc.communicate()
    kver = out.rstrip()
    return kver

'''
Get distribution name
'''
def get_dist_name():
    try:
        dist = list(platform.linux_distribution())

        # A little magic with Virtuozzo / Virtuozzo Linux
        if dist[0] == 'Virtuozzo Linux':
            for line in (open('/etc/os-release')):
                 m = re.match('^NAME\\s*=\\s*\"(.+)\"', line)
                 if m:
                     dist[0] = m.group(1)
                     break

        distname = dist[0].replace(' ', '-') + '-' + dist[1].split('.')[0]
        if distname.startswith("Ubuntu"):
            distname += "-" + dist[1].split('.')[1]


        # Check if distribution is supported
        if not get_api_result(APIURL + "/distros/" + distname):
            print("Your distribution (" + distname + ") is not supported.")
            print("For details, check " + BASEURL)
            sys.exit(1)

        return distname
    except Exception as e:
        print('Unable to get a distribution name')
        sys.exit(2)

'''
Send API request and return resulting JSON
'''
def get_api_result(api_url):
    try:
        j = urllib2.urlopen(api_url).read()
    except urllib2.HTTPError as e:
        if e.code == 404:
            return None
        print_warning(e)
        sys.exit(10)
    except urllib2.URLError as e:
        print_warning(e)
        sys.exit(11)
    return j

'''
Get a list of issues fixed by the 'latest_patch' patch.
latest_patch should be a patch url for RK API
'''
def get_CVEs(latest_patch, code=None):
    try:
        j = urllib2.urlopen(latest_patch).read()
    except urllib2.HTTPError as e:
        if e.code == 404:
            print("Failed to load information about available patches")
            sys.exit(5)
        print_warning(e)
        sys.exit(3)
    except urllib2.URLError as e:
        print_warning(e)
        sys.exit(4)

    if code:
        p_code = json.loads(j).get('code', False)
        if p_code != code:
            return []

    try:
        CVEs = json.loads(j).get('vulnerabilities', False)
        return CVEs
    except:
        print("Failed to get list of known vulnerabilities")
        sys.exit(6)

'''
Print detailed list of issues that exist in the running kernel and can be fixed by RK
'''
def dump_CVEs(CVEs):
    if len(CVEs) == 0:
        print('It looks that you have fixes for all vulnerabilities known to ReadyKernel!')
        return

    print('ReadyKernel can fix the following problems on your system:')

    i = 0
    for vuln in CVEs:
        i = i + 1
        print '#',i,' : ',
        try:
            cve_id = "[" + vuln['uri'].split("/")[-2] + "] "
        except:
            cve_id = ""
        if 'short_description' in vuln:
            print cve_id + str(vuln['short_description'])
        if 'url' in vuln:
            if vuln['url']:
                print "     " + str(vuln['url'])

'''
Print short summary about issues that exist in the running kernel and can be fixed by RK
'''
def dump_summary(CVEs, distro, code):
    if len(CVEs) == 0:
        print('It looks that you have fixes for all vulnerabilities known to ReadyKernel!')
        return

    url = BASEURL + "/patch/"+ distro + "/" + str(code)
    print("Your kernel has " + str(len(CVEs)) + " known issues, you may read about them here:")
    print(url)
    print("You may use ReadyKernel service to keep your server up to date!")

'''
Get code of currently installed RK patch
'''
def get_current_patch_name():
    proc = subprocess.Popen(['readykernel', 'info'],  stdout=subprocess.PIPE)
    for line in proc.stdout:
        if line.startswith("Patch name"):
            return line.replace("Patch name: ", "").rstrip()

    return ""

'''
Check if we have already installed some RK patch
and substitute issues fixed by it from our statistics
'''
def substitute_fixed_CVEs(CVEs, all_patches):
    code = get_current_patch_name()
    if not code:
        return CVEs

    new_CVEs = []
    for p in all_patches:
        fixed_CVEs = get_CVEs(p, code)
        if not fixed_CVEs:
            continue

        for f in CVEs:
            if not f in fixed_CVEs:
                new_CVEs.append(f)
        break

    return new_CVEs

'''
Check if current kernel reached EOM in RK
'''
def check_eom(eom_ts):
    eom_date  = eom_ts.split("T")[0]
    (y, m, d) = eom_date.split("-")

    now = datetime.datetime.now()

    d_eom = datetime.date(int(y), int(m), int(d))
    d_now = datetime.date(now.year, now.month, now.day)

    return d_eom < d_now

def dump_header(distname, kver):
    print("ReadyKernel-scan    v0.9    [%s]" % (BASEURL))
    print("            Keep your kernel up to date!             ")
    print("Detected kernel version: " + kver)
    print("Detected distribution: " + distname)


'''
Main part
'''

def parse_cmd_line():
    global cmdline
    parser = argparse.ArgumentParser(description='ReadyKernel Vulnerability Scanner.')
    parser.add_argument('--summary', action='store_true', help='only show summary for the found vulnerabilities')

    cmdline = parser.parse_args(sys.argv[1:])

if __name__ == '__main__':
    parse_cmd_line()

    # Check if distribution is supported
    try:
        distname = os.environ["TEST_RK_DIST_NAME"]
    except:
        distname = get_dist_name()

    try:
        kver = os.environ["TEST_RK_KVER"]
    except:
        kver = get_kernel_ver()

    kernel_url = APIURL + "/distros/" + distname + "/kernels/" + kver + "/"
    j = get_api_result(kernel_url)
    if not j:
        print("Kernel " + kver + " is not supported for your OS distribution (%s)." %(distname))
        print("For details, check " + BASEURL)
        sys.exit(7)

    dump_header(distname, kver)

    try:
        all_patches = json.loads(j).get('patch_set', False)
        eom = check_eom(json.loads(j).get('eom_date', False))
    except:
        print("Failed to load information about patches for the current kernel.")
        sys.exit(8)

    latest_patch = all_patches[0]
    CVEs = get_CVEs(latest_patch)

    if os.path.isfile("/sbin/readykernel"):
        CVEs= substitute_fixed_CVEs(CVEs, all_patches)

    if cmdline.summary:
        try:
            patch_desc = get_api_result(latest_patch)
            code = json.loads(patch_desc).get('code', False)
        except:
            print("Failed to get kernel code for the current kernel.")
            sys.exit(9)
        dump_summary(CVEs, distname, code)
    else:
        dump_CVEs(CVEs)

    if eom:
        print("\n !!! Your kernel was supported by ReadyKernel but reached EOM date !!!\n")
