#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# API module of cve-search. Returns queries in JSON format
#
# Software is free software released under the "Modified BSD license"
#

# Copyright (c) 2013-2016 	Alexandre Dulaunoy - a@foo.be
# Copyright (c) 2014-2017 	Pieter-Jan Moreels - pieterjan.moreels@gmail.com

# imports
import json, os, random, re, datetime

from io import StringIO
from bson import json_util
from flask import Flask, request, Response, render_template
from functools import wraps

try:
    from flask_wtf import FlaskForm
except ImportError:
    from flask_wtf import Form as FlaskForm
    
from wtforms import StringField, SelectField, SubmitField, HiddenField
from wtforms.validators import DataRequired

try:
    from flask_wtf.csrf import CSRFProtect
except ImportError:
    from flask_wtf.csrf import CsrfProtect as CSRFProtect

from ivulncheck_web.lib.config_web import Configuration as conf
from ivulncheck_web.lib import ivulncheck_api_layer as ial
from ivulncheck_web.lib.ivulncheck_api_layer import APIError

pageLength = conf.getPageLength()
cveSearchURL = conf.getCVESearchAPIHost() + ":" + str(conf.getCVESearchAPIPort())

app = Flask(__name__, static_folder=os.path.join(os.path.dirname(__file__), "static"), static_url_path='/static', template_folder=os.path.join(os.path.dirname(__file__), "templates"))
app.config['SECRET_KEY'] = str(random.getrandbits(256))

# ~ csrf = CSRFProtect()
# ~ csrf.init_app(app)

#########
# FORMS #
#########

class HostsForm(FlaskForm):
    r = HiddenField(label="r", id="r", default=0)
    pageLength = HiddenField(label="pageLength", id="pageLength", default=40)
    
    hostname = StringField(label="Hostname", id="hostname", default="all", validators=[DataRequired()])
    os_version = SelectField(label="OS Version", id="fixed", default="all", choices = [("all","all")], validators=[DataRequired()])
    
    submit = SubmitField(label="Search") 
    
def edit_hosts():
    os_versions = []
    filters = {'pageLength': pageLength, 'r': 0, 'hostname': "all", 'os_version': "all"}
    try:
        os_versions = ial.request_api('GET', "/ivulncheck_api/os_versions", filters)
    except APIError as ae:
        return render_template('error.html',status={'except':'api-error','info': ae.message}) 
    form = HostsForm()
    for ov in os_versions:
        form.os_version.choices.append((ov, ov))
    return form


class VulnForm(FlaskForm):
    
    r = HiddenField(label="r", id="r", default=0)
    hostname = HiddenField(label="Hostname", id="hostname", default="all")
    pageLength = HiddenField(label="pageLength", id="pageLength", default=40)
    
    cveid = StringField(label="CVE ID", id="cveid", default="all", validators=[DataRequired()])
    package = StringField(label="Package", id="package", default="all", validators=[DataRequired()])
    urgency = SelectField(label= "Urgency", id="urgency", default="all", choices = [("all", "all"), ("unknown","unknown"), ("negligible","negligible"), ("low", "low"), ("medium", "medium"), ("high", "high"), ("critical", "critical")], validators=[DataRequired()])
    remotely_exp = SelectField(label="Remotely Exploitable", id="remotely_exp", default="all", choices = [("all", "all"), ("unknown", "unknown"), ("yes", "yes"), ("no", "no")], validators=[DataRequired()])
    fixed = SelectField(label="Fixed", id="fixed", default="all", choices = [("all","all"), ("yes", "yes"), ("no", "no")], validators=[DataRequired()])
    src = SelectField(label="Source", id="src", default="all", choices = [("all", "all"), ("debsecan", "debsecan"), ("ubuntu", "ubuntu"), ("cve", "cve")], validators=[DataRequired()])
    submit = SubmitField(label="Search")    
    
class FixForm(FlaskForm):
    r = HiddenField(label="r", id="r", default=0)
    hostname = HiddenField(label="Hostname", id="hostname", default="all")
    pageLength = HiddenField(label="pageLength", id="pageLength", default=40)
    
    package = StringField(label="Package", id="package", default="all", validators=[DataRequired()])
    urgency = SelectField(label= "Urgency", id="urgency", default="all", choices = [("all", "all"), ("unknown","unknown"), ("negligible","negligible"), ("low", "low"), ("medium", "medium"), ("high", "high"), ("critical", "critical")], validators=[DataRequired()])
    remotely_exp = SelectField(label="Remotely Exploitable", id="remotely_exp", default="all", choices = [("all", "all"), ("unknown", "unknown"), ("yes", "yes"), ("no", "no")], validators=[DataRequired()])
    src = SelectField(label="Source", id="src", default="all", choices = [("all", "all"), ("debsecan", "debsecan"), ("ubuntu", "ubuntu"), ("cve", "cve")], validators=[DataRequired()])
    skipped = SelectField(label="Skipped", id="skipped", default="false", choices = [("false", "false"), ("true", "true")], validators=[DataRequired()])
    submit = SubmitField(label="Search") 

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)
    
##########
# ROUTES #
##########

# /
def index():
    filters = {}
    if request.method == 'GET':
        filters = {'pageLength': pageLength, 'r': 0, 'hostname': "all", 'os_version': "all"}
    elif request.method == 'POST':
        filters = getFilterSettingsFromPost()
        filters['pageLength'] = pageLength
    try:
        hosts = ial.request_api('POST', "/ivulncheck_api/hosts", filters)
        vulns_stats = ial.request_api('GET', "/ivulncheck_api/vulns/stats")
        fixes_stats = ial.request_api('GET', "/ivulncheck_api/fixes/stats")
    except APIError as ae:
        return render_template('error.html',status={'except':'api-error','info': ae.message})
    if not hosts:
        return render_template('error.html',status={'except':'hosts-found','info': ""})
    return render_template('index.html', hosts=hosts, v_stats=vulns_stats, f_stats=fixes_stats, filters=filters)

# /hosts
def hosts():
    form = edit_hosts()
    filters = {}
    if request.method == 'GET':
        filters = {'pageLength': pageLength, 'r': 0, 'hostname': "all", 'os_version': "all"}
    elif request.method == 'POST':
        filters = getFilterSettingsFromPost()
        filters['pageLength'] = pageLength
    try:
        hosts = ial.request_api('POST', "/ivulncheck_api/hosts", filters)
        os_versions = ial.request_api('GET', "/ivulncheck_api/os_versions")
        hosts_stats = ial.request_api('GET', "/ivulncheck_api/hosts/stats")
    except APIError as ae:
        return render_template('error.html',status={'except':'api-error','info': ae.message})
    if not hosts:
        return render_template('error.html',status={'except':'hosts-found','info': ""})
    return render_template('hosts.html', hosts=hosts, os_versions=os_versions, stats=hosts_stats, form=form)
    
# /host/<hostname>
def host(hostname):
    try:
        host_stats = ial.request_api('GET', "/ivulncheck_api/host/stats/" + hostname)
    except APIError as ae:
        return render_template('error.html',status={'except':'api-error','info': ae.message})
    if host_stats == {"fixed_packages": 0, "nb_fixes": 0, "nb_vulns": 0, "up_packages": 0, "vuln_packages": 0}:
        return render_template('error.html',status={'except':'host-found','info': {'host':hostname}})
    return render_template('host.html', hostname=hostname, stats=host_stats)

# /host/vulns/<hostname>
def host_vulns(hostname):
    form = VulnForm()
    filters = {}
    if request.method == 'GET':
        filters = {'cveSearchURL': cveSearchURL, 'pageLength': pageLength, 'r': 0, 'hostname': hostname, 'cveid': "all", 'package': "all", 'urgency': "all", 'remotely_exp': "all", 'os_version': "all", 'fixed': "all", 'src': "all"}
    elif request.method == 'POST':
        filters = getFilterSettingsFromPost()
        filters['pageLength'] = pageLength
        filters['cveSearchURL'] = cveSearchURL
    try:
        host_vulns = ial.request_api('POST', "/ivulncheck_api/vulns", filters)
    except APIError as ae:
        return render_template('error.html',status={'except':'api-error','info': ae.message})
    if not host_vulns:
        return render_template('error.html',status={'except':'host-vulns-found','info':{'host':hostname}})
    return render_template('host_vulns.html', hostname=hostname, vulns=host_vulns, form=form)
    
# /host/fixes/<hostname>
def host_fixes(hostname):
    form = FixForm()
    filters = {}
    try:
        os_versions = ial.request_api('GET', "/ivulncheck_api/os_versions", filters)
    except APIError as ae:
        return render_template('error.html',status={'except':'api-error','info': ae.message}) 
    if request.method == 'GET':
        filters = {'cveSearchURL': cveSearchURL, 'pageLength': pageLength, 'r': 0, 'hostname': hostname, 'package': "all", 'urgency': "all", 'remotely_exp': "all", 'os_version': "all", 'src': "all", 'skipped': False}
    elif request.method == 'POST':
        filters = getFilterSettingsFromPost()
        filters['pageLength'] = pageLength
        filters['cveSearchURL'] = cveSearchURL
        filters['hostname'] = hostname
        
    try:
        host_fixes = ial.request_api('POST', "/ivulncheck_api/fixes", filters)
    except APIError as ae:
        return render_template('error.html',status={'except':'api-error','info': ae.message})
    if not host_fixes and filters == {'cveSearchURL': cveSearchURL, 'pageLength': pageLength, 'r': 0, 'hostname': hostname, 'package': "all", 'urgency': "all", 'remotely_exp': "all", 'os_version': "all", 'src': "all", 'skipped': False}:
        return render_template('error.html',status={'except':'host-fixes-found','info':{'host':hostname}})
    return render_template('host_fixes.html', hostname=hostname, fixes=host_fixes, os_versions=os_versions, form=form)
      
# /vulns
def vulns():
    form = VulnForm()
    filters = {}
    if request.method == 'GET':
        filters = {'cveSearchURL': cveSearchURL, 'pageLength': pageLength, 'r': 0, 'hostname': "all", 'cveid': "all", 'package': "all", 'urgency': "all", 'remotely_exp': "all", 'fixed': "all", 'src': "all"}
    elif request.method == 'POST':
        filters = getFilterSettingsFromPost()
        filters['pageLength'] = pageLength
        filters['cveSearchURL'] = cveSearchURL
    try:
        vulns = ial.request_api('POST', "/ivulncheck_api/vulns", filters)
        vulns_stats = ial.request_api('GET', "/ivulncheck_api/vulns/stats")
    except APIError as ae:
        return render_template('error.html',status={'except':'api-error','info': ae.message})
    if not vulns and filters == {'cveSearchURL': cveSearchURL, 'pageLength': pageLength, 'r': 0, 'hostname': "all", 'cveid': "all", 'package': "all", 'urgency': "all", 'remotely_exp': "all", 'fixed': "all", 'src': "all"}:
        return render_template('error.html',status={'except':'vulns-found','info': ""})
    return render_template('vulns.html', vulns=vulns, stats=vulns_stats, form=form, cveSearchURL=cveSearchURL)
    
# /fixes
def fixes():
    form = FixForm()
    filters = {}
    # ~ try:
        # ~ os_versions = ial.request_api('GET', "/ivulncheck_api/os_versions", filters)
    # ~ except APIError as ae:
        # ~ return render_template('error.html',status={'except':'api-error','info': ae.message}) 
    if request.method == 'GET':
        filters = {'cveSearchURL': cveSearchURL, 'pageLength': pageLength, 'r': 0, 'hostname': "all", 'package': "all", 'urgency': "all", 'remotely_exp': "all", 'os_version': "all", 'src': "all", 'skipped': False}
    elif request.method == 'POST':
        filters = getFilterSettingsFromPost()
        filters['pageLength'] = pageLength
        filters['cveSearchURL'] = cveSearchURL
    try:
        fixes = ial.request_api('POST', "/ivulncheck_api/fixes", filters)
        fixes_stats = ial.request_api('GET', "/ivulncheck_api/fixes/stats")
    except APIError as ae:
        return render_template('error.html',status={'except':'api-error','info': ae.message})
    if not fixes and filters == {'cveSearchURL': cveSearchURL, 'pageLength': pageLength, 'r': 0, 'hostname': "all", 'package': "all", 'urgency': "all", 'remotely_exp': "all", 'os_version': "all", 'src': "all", 'skipped': False}:
        return render_template('error.html',status={'except':'fixes-found','info': ""})
    return render_template('fixes.html', fixes=fixes, stats=fixes_stats, form=form, cveSearchURL=cveSearchURL)
    
##################
# Error Messages #
##################
def page_not_found(e):
    return render_template('404.html'), 404

routes = [{'r': '/',                               'm': ['GET'],  'f': index},
              
          {'r': '/hosts',                          'm': ['GET'],  'f': hosts},
          {'r': '/hosts',                          'm': ['POST'],  'f': hosts},
          {'r': '/host/<hostname>',                'm': ['GET'],  'f': host},
          {'r': '/host/fixes/<hostname>',          'm': ['GET'],  'f': host_fixes},
          {'r': '/host/fixes/<hostname>',          'm': ['POST'],  'f': host_fixes},
          {'r': '/host/vulns/<hostname>',          'm': ['GET'],  'f': host_vulns},
          {'r': '/host/vulns/<hostname>',          'm': ['POST'],  'f': host_vulns},
          {'r': '/fixes',                          'm': ['GET'],  'f': fixes},
          {'r': '/fixes',                          'm': ['POST'],  'f': fixes},
          {'r': '/vulns',                          'm': ['GET'],  'f': vulns},
          {'r': '/vulns',                          'm': ['POST'],  'f': vulns}]
    
error_handlers = [{'e': 404, 'f': page_not_found}]

for route   in routes:             addRoute(route)
for handler in error_handlers:     app.register_error_handler(handler['e'], handler['f'])

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