# API of ivulncheck
#

import json, logging, os, random, re, datetime, sys

from io import StringIO

from bson               import json_util

from flask              import Flask, request, Response, render_template
from functools          import wraps

from ivulncheck_api.lib.database_layer import DatabaseLayer
from ivulncheck_api.bin import insert as ins
from ivulncheck_api.bin import check as chk
from ivulncheck_api.bin import download as dl

from ivulncheck_api.lib.config_api import Configuration as conf

import logging
logging.basicConfig(filename=conf.getLogfile(), level=logging.INFO)

class APIError(Exception):
    def __init__(self, message, status=500):
        self.message = message
        self.status  = status


app = Flask(__name__, template_folder=os.path.join(os.path.dirname(__file__), "templates"))

app.config['SECRET_KEY'] = str(random.getrandbits(256))

def init_db():
    # init database
    db = DatabaseLayer()
    try:
        db.init()
        return db
    except Exception as e:
        logging.info("[%s] Database unreachable : %s", datetime.datetime.today().isoformat(), str(e))
        raise(APIError("Database unreachable!"))
    

def addRoute(route):
    app.add_url_rule(route['r'], view_func=route['f'], methods=route['m'])

def getFilterSettingsFromPost():
    filters = dict(request.form)
    filters = {x: filters[x][0] for x in filters.keys()}
    return(filters)
    

#############
# Decorator #
#############

def api(funct):
    @wraps(funct)
    def api_wrapper(*args, **kwargs):
        data = error = None
        # Get data (and possibly errors)
        try:
            data = funct(*args, **kwargs)
        except APIError as ae:
            logging.info("[%s] APIError : [%d] : %s", datetime.datetime.today().isoformat(), ae.status, ae.message)
            return Response(json.dumps({'error': "APIError", 'status': ae.status, 'message': ae.message}), mimetype='application/json'), ae.status
        except Exception as e:
            logging.info("[%s] Error :  %s", datetime.datetime.today().isoformat(), str(e))
            return Response(json.dumps({'error': "Error", 'status': 500, 'message': str(e)}), mimetype='application/json'), 500
        # Check if data should be returned as html or data
        try:
            returnType = 'application/json'
            
            if type(data) is tuple:
                data = list(data)
                data[0] = json.dumps(data[0], indent=None, sort_keys=True, default=json_util.default)
            else:
                data = (json.dumps(data, indent=None, sort_keys=True, default=json_util.default), 200)
            
            return Response(data[0], mimetype=returnType), data[1]
        except Exception as e:
            print(e)
            pass
        return data
    return api_wrapper
    

#############
# FUNCTIONS #
#############

def generate_hosts_query(f):
    query = []
    if f['hostname'] != "all":
        query.append({'id': f['hostname']})
    if f['os_version'] != "all":
        query.append({'os_version': f['os_version']})
    return query

def hosts_filter_logic(filters, limit=None):
    db = init_db()
    limit = limit if limit else int(filters['pageLength'])
    query = generate_hosts_query(filters)
    return db.getHosts(limit=limit, skip=int(filters['r']), query=query)

def pkgs_filter_logic(filters, limit=None, hide_hosts=False):
    db = init_db()
    limit = limit if limit else int(filters['pageLength'])
    return db.getPackages(limit=limit, skip=int(filters['r']), hide_hosts=hide_hosts)

def generate_vulns_query(f):
    query = []
    if f['cveid'] != "all":
        query.append({'id': f['cveid']})
    if f['package'] != "all":
        if bool(re.match(r'.*\*.*', f['package'])) == True:
            f['package'].replace('*', '')
            query.append({'bin_package': {'$regex': ".*" + f['package'] + ".*"}})
        else:
            query.append({'bin_package': f['package']})
    if f['urgency'] != "all":
        query.append({'urgency': f['urgency']})
    if f['remotely_exp'] == "yes":
        query.append({'remotely_exp': True})
    elif f['remotely_exp'] == "no":
        query.append({'remotely_exp': False})
    elif f['remotely_exp'] == "unknown":
        query.append({'remotely_exp': None})
    if f['fixed'] == "yes":
        query.append({'fixed': {'$ne': []}})
    elif f['fixed'] == "no":
        query.append({'fixed': []})
    if f['src'] != "all":
        query.append({'src': f['src']})
        
    return query

def vulns_filter_logic(filters, limit=None):
    db = init_db()
    query = generate_vulns_query(filters)
    limit = limit if limit else int(filters['pageLength'])
    if filters['hostname'] != "all":
        return db.getHostVulns(hostname=filters['hostname'], limit=limit, skip=int(filters['r']), query=query)
    return db.getVulns(limit=limit, skip=int(filters['r']), query=query)
    
def generate_fixes_query(f):
    query = []
    if f['hostname'] != "all":
        query.append({'hosts': f['hostname']})
    if f['os_version'] != "all":
        query.append({'os_version': f['os_version']})
    if f['package'] != "all":
        if bool(re.match(r'.*\*.*', f['package'])) == True:
            query.append({'bin_package': {'$regex': ".*" + f['package'] + ".*"}})
        else:
            query.append({'bin_package': f['package']})
    if f['urgency'] != "all":
        query.append({'urgency': f['urgency']})
    if f['remotely_exp'] == "yes":
        query.append({'remotely_exp': True})
    elif f['remotely_exp'] == "no":
        query.append({'remotely_exp': False})
    elif f['remotely_exp'] == "unknown":
        query.append({'remotely_exp': None})
    query.append({'skipped': f['skipped']})
    if f['src'] != "all":
        query.append({'src': f['src']})
        
    return query

def fixes_filter_logic(filters, limit=None, hide_hosts=False):
    db = init_db()
    if filters['skipped'] == "False" or filters['skipped'] == "false":
        filters['skipped'] = False
    elif filters['skipped'] == "True" or filters['skipped'] == "true":
        filters['skipped'] = True
    query = generate_fixes_query(filters)
    limit = limit if limit else int(filters['pageLength'])
    return db.getFixes(limit=limit, skip=int(filters['r']), query=query, hide_hosts=hide_hosts)


##########
# ROUTES #
##########

# /api/vulns
@api
def api_vulns():
    db = init_db()
    vulns = db.getVulns()
    if not vulns: raise(APIError('vulns not found', 404))
    return vulns  
  
@api
def api_vulns_post():
    filters = getFilterSettingsFromPost()
    vulns = vulns_filter_logic(filters)
    if not vulns: raise(APIError('vulns not found', 404))
    return vulns

# /api/vuln/<vulnid>
@api
def api_vuln(vulnid):
    db = init_db()
    vuln = db.getVulns(query={'id': vulnid})
    if not vuln: raise(APIError('vuln not found', 404))
    return vuln

# /api/pkgs
@api
def api_pkgs():
    db = init_db()
    pkgs = db.getPackages()
    if not pkgs: raise(APIError('packages not found', 404))
    return pkgs
    
@api
def api_pkgs_post():
    filters = getFilterSettingsFromPost()
    hide_hosts = False
    if filters['hostname'] != "all":
        hide_hosts = True
    pkgs = pkgs_filter_logic(filters, hide_hosts=hide_hosts)
    if not pkgs: raise(APIError('packages not found', 404))
    return pkgs

# /api/pkg/<pkgvers>
@api
def api_pkg(pkgvers):
    db = init_db()
    try:
        if re.match(r':amd64', pkgvers) == False:
            pkg = pkgvers.split(':')[0]
            vers = pkgvers.split(':')[1]
        else:
            pkg = pkgvers
            vers = ""
    except IndexError:
        pkg = pkgvers
        vers = ""
    pkg = db.getPackage(pkg, vers)
    if not pkg: raise(APIError('package not found', 404))
    return pkg

# /api/pkgsearch/<pkgvers>
@api
def api_pkgsearch(pkgvers):
    db = init_db()
    pkg = db.findPackage(pkgvers)
    if not pkg: raise(APIError('package not found', 404))
    return pkg

# /api/hosts
@api
def api_hosts():
    db = init_db()
    hosts = db.getHosts()
    if not hosts: raise(APIError('hosts not found', 404))
    return hosts
    
@api
def api_hosts_post():
    filters = getFilterSettingsFromPost()
    hosts = hosts_filter_logic(filters)
    if not hosts: raise(APIError('hosts not found', 404))
    return hosts
    
# /api/hosts/stats
@api
def api_hosts_stats():
    db = init_db()
    hosts_stats = db.getHostsStats()
    if not hosts_stats: raise(APIError('hosts stats not found', 404))
    return hosts_stats

# /api/host/<hostname>
@api
def api_host(hostname):
    db = init_db()
    host = db.getHosts(query={'id': hostname})
    if not host: raise(APIError('host not found', 404))
    return host
    
# /api/host/stats/<hostname>
@api
def api_host_stats(hostname):
    db = init_db()
    host_stats = db.getHostStats(hostname)
    if not host_stats: raise(APIError('host s stats not found', 404))
    return host_stats

# /api/host/pkgs/<hostname>
@api
def api_host_pkgs(hostname):
    db = init_db()
    host_pkgs = db.getPackages(query={'hosts': hostname}, hide_hosts=True)
    if not host_pkgs: raise(APIError('host s packages not found', 404))
    return host_pkgs

# /api/host/vulns/<hostname>
@api
def api_host_vulns(hostname):
    db = init_db()
    host_vulns = db.getHostVulns(hostname)
    if not host_vulns: raise(APIError('host s vulns not found', 404))
    return host_vulns

# /api/host/fixes/<hostname>
@api
def api_host_fixes(hostname):
    db = init_db()
    host_fixes = db.getFixes(query={'hosts': hostname}, hide_hosts=True)
    if not host_fixes: raise(APIError('host s fixes not found', 404))
    return host_fixes
    
# /api/fixes
@api
def api_fixes():
    db = init_db()
    fixes = db.getFixes()
    if not fixes: raise(APIError('fixes not found', 404))
    return fixes

@api
def api_fixes_post():
    filters = getFilterSettingsFromPost()
    hide_hosts = False
    if filters['hostname'] != "all":
        hide_hosts = True
    fixes = fixes_filter_logic(filters, hide_hosts=hide_hosts)
    if not fixes: raise(APIError('fixes not found', 404))
    return fixes
    
# /api/vulns/stats
@api
def api_vulns_stats():
    db = init_db()
    vulns_stats = db.getVulnsStats()
    if not vulns_stats: raise(APIError('vulns stats not found', 404))
    return vulns_stats
    
# /api/fixes/stats
@api
def api_fixes_stats():
    db = init_db()
    fixes_stats = db.getFixesStats()
    if not fixes_stats: raise(APIError('fixes stats not found', 404))
    return fixes_stats

# /api/os_versions
@api
def api_os_versions():
    db = init_db()
    os_versions = db.getOSVersions()
    if not os_versions: raise(APIError('os_versions not found', 404))
    return os_versions
    
# /api/suites
@api
def api_suites():
    db = init_db()
    suites = db.getSuites()
    if not suites: raise(APIError('suites not found', 404))
    return suites
    
# /api/releases
@api
def api_releases():
    db = init_db()
    releases = db.getReleases()
    if not releases: raise(APIError('releases not found', 404))
    return releases
    
# /api/cmd/purge
@api
def api_cmd_purge(collection):
    db = init_db()
    results = []
    logging.info("[%s] Trying to purge collection %s", datetime.datetime.today().isoformat(), collection)
    collections = ["hosts", "packages", "products", "vulns", "fixes", "pkg2cpe"]
    if collection in collections:
        results.append(db.purgeCollection(collection))
    elif collection == "all":
        results = db.purgeAll()
    else:
        raise(APIError('no such collection to purge : ' + collection, 404))
        
    for result in results:
        if result.acknowledged == True:
            return {"action": "Purge collection " + collection, "message":  collection + " purged."}
        else:
            return {"action": "Purge collection " + collection, "message":  collection + "not purged."}

# /api/cmd/purge/host/<hostname>
@api
def api_cmd_purge_host(hostname):
    db = init_db()
    result = None
    logging.info("[%s] Trying to purge host %s", datetime.datetime.today().isoformat(), hostname)

    if db.getHosts(query={'id': hostname}) is not None:
        results = db.purgeHost(hostname)
    else:
        raise(APIError('no such host to purge : ' + hostname, 404))
    for result in results:
        if result.acknowledged == False:
            return {"action": "Purge host " + hostname, "message":  hostname + "not completely purged."}
    return {"action": "Purge host " + hostname, "message":  hostname + " purged."}
        
        
# /api/cmd/check
@api
def api_cmd_check(action):
    db = init_db()
    if action == "pkg2cpe":
        (nb_mod, nb_err) = chk.pkg2cpe(db)
        return {"action": "Check CPE from CVE-Search", "message":  str(nb_mod) + " indexes pkg2cpe added and " + str(nb_err) + " errors."}
    elif action == "debsecan_vulns":
        (nb_mod, nb_err) = chk.debsecan_vulns(db)
        return {"action": "Check vulnerabilities from debsecan", "message":  str(nb_mod) + " vulnerabilities added and " + str(nb_err) + " errors."}
    elif action == "products_vulns":
        (nb_mod, nb_err) = chk.products_vulns(db)
        return {"action": "Check vulnerabilities for products from CVE-Search", "message":  str(nb_mod) + " vulnerabilities added."}
    elif action == "ubuntu_vulns":
        (nb_mod, nb_err) = chk.ubuntu_vulns(db)
        return {"action": "Check vulnerabilities from Ubuntu-CVE-Tracker", "message":  str(nb_mod) + " vulnerabilities added."}
    elif action == "debsecan_fixes":
        (nb_mod, nb_err) = chk.debsecan_fixes(db)
        return {"action": "Check fixes from debsecan fixed tags", "message":  str(nb_mod) + " fixes added."}
    elif action == "ubuntu_fixes":
        (nb_mod, nb_err) = chk.ubuntu_fixes(db)
        return {"action": "Check fixes from Ubuntu-CVE-Tracker", "message":  str(nb_mod) + " fixes added."}
    elif action == "madison_fixes":
        (nb_deb, nb_ub, nb_mod, nb_err) = chk.madison_fixes(db)
        return {"action": "Check fixes from Madison", "message":  str(nb_deb) + " fixes for debian found, " + str(nb_ub) + " fixes found for ubuntu and " + str(nb_mod) + " added."}
    elif action == "cve_vulns":
        chk.pkg2cpe(db)
        (nb_mod, nb_err) = chk.cve_vulns(db)
        return {"action": "Check vulnerabilities for packages from CVE-Search", "message":  str(nb_mod) + " vulnerabilities added."}
    elif action == "deb2uburg":
        (nb_mod, nb_err) = chk.debian2ubuntu_urgency(db)
        return {"action": "Complete Ubuntu Fixes with Debian Informations", "message":  str(nb_mod) + " ubuntu fixes completed with debian informations."}
    else:
        raise(APIError('no such action to check : ' + action, 404))
        
# /api/cmd/calcstats
@api
def api_cmd_calc_stats(hostname):
    db = init_db()
    result = db.calculateHostStats(hostname)
    if result == True:
        return {"action": "Calc Stats", "message":  "Host Stats for " + hostname + " inserted."}
    else:
        raise(APIError('An error occured during host stats calulation for ' + hostname + '.' , 404))
        
# /api/cmd/fake_repos
@api
def api_cmd_fake_repos():
    db = init_db()
    result = dl.download_fake_repos(db)
    if result == True:
        return {"action": "Fake Repos", "message":  "Fake repos downloaded."}
    else:
        raise(APIError('An error occured during fake repos download.' , 404))

# /api/cmd/uct
@api
def api_cmd_uct():
    db = init_db()
    result = dl.update_ubuntu_cve_tracker()
    if result == True:
        return {"action": "Ubuntu-CVE-Tracker", "message":  "ubuntu-cve-tracker updated."}
    else:
        raise(APIError('An error occured during ubuntu-cve-tracker update.' , 404))

# /api/cmd/skip
@api
def api_cmd_skip(pkgvers):
    db = init_db()
    result = None
    if re.match(r'[\w\-\:]+\:.+', pkgvers):
        (pkg, vers) = pkgvers.split(':', 2)
        result = db.skipFix(pkg, vers)
        if result == True:
            return {"action": "Skip fix", "message":  "fix for " + pkg + " package and version " + version + " has been skipped."}
        else:
            raise(APIError('fix for ' + pkg + ' package and version ' + vers + ' not skipped.' , 404))
            
    else:
        raise(APIError('pkgvers must be <pkgname>:<pkgvers>.', 404))
        
# /api/cmd/unskip
@api
def api_cmd_unskip(pkgvers):
    db = init_db()
    result = None
    if re.match(r'[\w\-\:]+\:.+', pkgvers):
        (pkg, vers) = pkgvers.split(':', 2)
        result = db.unskipFix(pkg, vers)
        if result == True:
            return {"action": "Unskip fix", "message":  "fix for " + pkg + " package and version " + vers + " has been unskipped."}
        else:
            raise(APIError('fix for ' + pkg + ' package and version ' + vers + ' not unskipped.', 404))
            
    else:
        raise(APIError('pkgvers must be <pkgname>:<pkgvers>.', 404))


# /api/insert/host
@api
def api_insert_host():
    db = init_db()
    # testing JSON format
    try:
        req_data = json.load(StringIO(request.data.decode("utf-8")))
        # Testing JSON Data
        try : 
            req_data['hostname']
            req_data['ipv4']
            req_data['osfamily']
            req_data['ostype']
            req_data['osversion']
            req_data['packages']
        except KeyError:
            raise(APIError('inserhost impossible : wrong JSON data', 400))
    except ValueError:
        raise(APIError('inserhost impossible : data not in JSON type', 400))
    (resulthost, nb_mod_pkg, nb_err_pkg, nb_mod_fix, nb_err_fix) = ins.insert_host(db, req_data)
    if resulthost == True:
        if nb_mod_fix != 0 or nb_err_fix != 0:
            return {"action": "Insert Host " + req_data['hostname'], "message": "Host " + req_data['hostname'] + " inserted and " + str(nb_mod_pkg) + " packages added with " + str(nb_err_pkg) + " errors and " + str(nb_mod_fix) + " ubuntu fixes with " + str(nb_err_fix) + " errors."}
        return {"action": "Insert Host " + req_data['hostname'], "message": "Host " + req_data['hostname'] + " inserted and " + str(nb_mod_pkg) + " packages added with " + str(nb_err_pkg) + " errors."}
    else:
        return {"action": "Insert Host " + req_data['hostname'], "message": "Host " + req_data['hostname'] + " not inserted check log file (probably already exists.)"} 

routes = [{'r': '/vulns',                           'm': ['GET'], 'f': api_vulns},
          {'r': '/vulns',                           'm': ['POST'], 'f': api_vulns_post},
          {'r': '/vuln/<vulnid>',                   'm': ['GET'], 'f': api_vuln},
          {'r': '/vulns/stats',                     'm': ['GET'], 'f': api_vulns_stats},
          
          {'r': '/pkgs',                            'm': ['GET'], 'f': api_pkgs},
          {'r': '/pkgs',                            'm': ['POST'], 'f': api_pkgs_post},
          {'r': '/pkg/<pkgvers>',                   'm': ['GET'], 'f': api_pkg},
          
          {'r': '/hosts',                           'm': ['GET'], 'f': api_hosts},
          {'r': '/hosts',                           'm': ['POST'], 'f': api_hosts_post},
          {'r': '/hosts/stats',                     'm': ['GET'], 'f': api_hosts_stats},
          {'r': '/host/<hostname>',                 'm': ['GET'], 'f': api_host},
          {'r': '/host/stats/<hostname>',           'm': ['GET'], 'f': api_host_stats},
          {'r': '/host/vulns/<hostname>',           'm': ['GET'], 'f': api_host_vulns},
          {'r': '/host/pkgs/<hostname>',            'm': ['GET'], 'f': api_host_pkgs},
          {'r': '/host/fixes/<hostname>',           'm': ['GET'], 'f': api_host_fixes},
          
          {'r': '/fixes',                           'm': ['GET'], 'f': api_fixes},
          {'r': '/fixes',                           'm': ['POST'], 'f': api_fixes_post},
          {'r': '/fixes/stats',                     'm': ['GET'], 'f': api_fixes_stats},
          
          {'r': '/os_versions',                          'm': ['GET'], 'f': api_os_versions},
          {'r': '/suites',                          'm': ['GET'], 'f': api_suites},
          {'r': '/releases',                        'm': ['GET'], 'f': api_releases},
          
          {'r': '/cmd/purge/<collection>',          'm': ['GET'], 'f': api_cmd_purge},
          {'r': '/cmd/purge/host/<hostname>',       'm': ['GET'], 'f': api_cmd_purge_host},
          {'r': '/cmd/fake_repos',                  'm': ['GET'], 'f': api_cmd_fake_repos},
          {'r': '/cmd/uct',                         'm': ['GET'], 'f': api_cmd_uct},
          {'r': '/cmd/check/<action>',              'm': ['GET'], 'f': api_cmd_check},
          {'r': '/cmd/calcstats/<hostname>',        'm': ['GET'], 'f': api_cmd_calc_stats},
          {'r': '/cmd/skip/<pkgvers>',              'm': ['GET'], 'f': api_cmd_skip},
          {'r': '/cmd/unskip/<pkgvers>',            'm': ['GET'], 'f': api_cmd_unskip},
          
          {'r': '/insert/host',                     'm': ['PUT'], 'f': api_insert_host}]
    

for route   in routes:             addRoute(route)

if __name__ == '__main__':
    app.run()
