#Global Vulnerability Database Project
#import reuired libraries
from flask import Flask, render_template, request, redirect, url_for, g, session
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import time,os
import sqlite3
import requests
from bs4 import BeautifulSoup
import hashlib
import random,string
import jsonify
import datetime
import smtplib
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText

#sending email with random password
def send_recovery_email(username , email ):
    #generate random password including uppercase lowercase and numbers
    password = ''.join(random.choices(string.ascii_uppercase +string.ascii_lowercase + string.digits, k = 8))
    salt = os.urandom(32)
    key = hashlib.pbkdf2_hmac('sha512', password.encode(), salt, 10000)
    storage = salt + key #passwors hash to be stored
    conn = sqlite3.connect('vulnerabilityDB.sqlite')
    cur = conn.cursor()
    #update hash in database
    cur.execute('''UPDATE adminhash SET hash = ? WHERE username = ? ''', (storage, username))
    conn.commit()
    conn.close()

    #send email

    mail_content = "Password Recovery:\nDear {}\nYour New password is:\n\t{}\nPlease login to your account and change your password.\nThank You".format(username.upper(),password)

    # The mail addresses and password
    sender_address = "abc@xyz.com" #Add sender Address Here
    sender_pass = "password" #Add sender password Here
    receiver_address = email #Add receiver/admin Address Here
    # Setup the MIME
    message = MIMEMultipart()
    message['From'] = sender_address
    message['To'] = receiver_address

    message['Subject'] = 'Password Reset'  # The subject line
    # The body and the attachments for the mail
    message.attach(MIMEText(mail_content, 'plain'))

    # Create SMTP session for sending the mail
    session = smtplib.SMTP('smtp.gmail.com', 587)  # use gmail with port
    session.starttls()  # enable security
    session.login(sender_address, sender_pass)  # login with mail_id and password
    text = message.as_string()
    session.sendmail(sender_address, receiver_address, text)
    session.quit()
    print('Mail Sent')






    return

#get list of cve for given keyword using selenium and beautiful soup
def get_cve_list(keyword):
    PATH = 'C:\Program Files (x86)\chromedriver.exe'
    #uncomment these 3 lines to avoid opening of browser
    #option = webdriver.ChromeOptions()
    #option.add_argument('headless')
    #driver = webdriver.Chrome(PATH, options=option)
    driver = webdriver.Chrome(PATH)
    driver.get("https://www.cvedetails.com/google-search-results.php?q={}&sa=Search".format(keyword))

    cve_vulnerability_list = []
    cve_ids_no = []
    for i in range(1, 11):
        #open and collect data from different pages of website
        elements = WebDriverWait(driver, 10).until(
            EC.element_to_be_clickable((By.XPATH, '/html/body/table[1]/tbody/tr[2]'
                                                  '/td[2]/div/div/div/div/div/div/'
                                                  'div[5]/div[2]/div/div/div[2]/'
                                                  'div/div[' + str(i) + ']')))
        elements.click()
        window_after = driver.window_handles[0]
        driver.switch_to.window(window_after)
        time.sleep(2)
        #search element using class name
        search = driver.find_elements_by_class_name("gs-per-result-labels")
        for link in search:
            link_check = link.get_attribute("url")
            if str("vulnerability-list") in str(link_check) or str("microsoft-bulletin") in str(link_check):
                cve_vulnerability_list.append(link_check)
            elif str("/cve/CVE-") in str(link_check):
                cve_id_no = link_check.split('/')[-2]
                if cve_id_no > "CVE-2010":
                    page = requests.get(str(link_check))
                    soup = BeautifulSoup(page.content, 'html.parser')
                    desc = soup.find_all('div', class_='cvedetailssummary')[0].text.split('\n')[1].strip()
                    cve_ids_no.append((cve_id_no,desc)) # Append cve id and description as tuple to a list

    driver.quit() #close web browser
    cve_vulnerability_list = list(dict.fromkeys(cve_vulnerability_list)) #delete duplicated
    #find cve id and description from a list of links
    for link in cve_vulnerability_list:
        page = requests.get(link)
        soup = BeautifulSoup(page.content, 'html.parser')
        for terms, desc in zip(soup.find_all('tr', class_='srrowns'), soup.find_all('td', class_='cvesummarylong')):
            for all_id in terms.find_all('a'):
                cve_id_no = all_id.getText()
                if "CVE" in cve_id_no:
                    if cve_id_no > "CVE-2010":
                        cve_ids_no.append((cve_id_no, desc.text.strip()))

    cve_ids_no = list(dict.fromkeys(cve_ids_no))  #delete duplicated
    return cve_ids_no

#update database
def custom_update():
    #connect to sql database
    conn = sqlite3.connect('vulnerabilityDB.sqlite')
    cur = conn.cursor()
    #get dictionary of keywords
    my_file = open("keyword.txt", 'r')
    content = my_file.read()
    my_file.close()
    keyword_list = content.split('\n')
    cve_ids_no = [('','')]
    for keyword in keyword_list:
        try:
            keyword = keyword[:-1] #reduce , at the end of keyword
            cve_ids_no = get_cve_list(keyword) #get list of cve and their description from get_cve_list
        except:
            # No record found for this keyword so add this keyword to blacklist
            my_file = open("black_listed_keys.txt", "a")
            my_file.write("\n" + keyword + ",")
            my_file.close()
        for cve_id ,des in cve_ids_no:
            # check if this cve exist in database
            query = '''SELECT CVE_ID,Description FROM vulnerabilityDB WHERE CVE_ID LIKE "{}" '''.format(cve_id)
            cur.execute(query)
            result = cur.fetchone()
            if result == None:
                # This cve doesn't exist in database so add it
                cur.execute('''INSERT INTO vulnerabilityDB (CVE_ID, Description, Keyword) VALUES (?,?,?)''', (cve_id, des,keyword.lower()))
                conn.commit()
            else:
                # This cve exist in database
                # Check if the cve is labeled with keyword
                query = '''SELECT Keyword FROM vulnerabilityDB WHERE CVE_ID LIKE "{}" '''.format(cve_id)
                cur.execute(query)
                result2 = cur.fetchone()
                existing_keyword = result2[0]
                if bool(existing_keyword):
                    # yes cve is labeled, now check if it is already labeled with current keyword
                    if keyword not in existing_keyword:
                        # append the keyword to existing keywords
                        new_keyword = keyword + ',' + existing_keyword
                        new_keyword = new_keyword.lower()
                        # label cve with list of new keywords
                        cur.execute('''UPDATE vulnerabilityDB SET Description = ?, Keyword = ? WHERE CVE_ID = ? ''', (des, new_keyword, cve_id))
                        conn.commit()
                    else:
                        pass # do nothing as it is labeled with given keyword
                else:
                    # cve is not labeled previously, just label it
                    keyword.lower()
                    cur.execute('''UPDATE vulnerabilityDB SET Description = ? , Keyword = ? WHERE CVE_ID = ? ''', (des, keyword, cve_id))
                    conn.commit()
    conn.close()
    return

app = Flask(__name__)
app.secret_key = os.urandom(24)

# Load home page for both of given urls
@app.route('/', methods = ['GET','POST'])
@app.route('/home', methods = ['GET','POST'])
def home():
    return render_template('home.html')

# load the search page and get keyword from user
@app.route('/search/', methods = ['GET','POST'])
def search():
    if request.method == 'POST':
        form = request.form
        keyword = form['keyword']
        return redirect(url_for('search_result', keyword = keyword))
    return render_template('search.html')

@app.route('/search/<string:keyword>')
def search_result(keyword):
    # check if the given keyword is black listed
    my_file = open("black_listed_keys.txt", 'r')
    black_keys = my_file.read()
    my_file.close()
    black_keys_list = black_keys.split('\n')
    if keyword not in black_keys_list:
        # as keyword is not blacklisted search it in database
        conn = sqlite3.connect('vulnerabilityDB.sqlite')
        cur = conn.cursor()
        query = "SELECT CVE_ID,Description FROM vulnerabilityDB WHERE Keyword LIKE '%{}%' OR CVE_ID LIKE '{}' ".format(keyword, keyword)
        cur.execute(query)
        result = cur.fetchall()
        if bool(result):
            #keyword found in database
            result = sorted(result, reverse=True)
            conn.close()
            return render_template('search_result.html',searchresult = result)
        else:
            try:
                # get list of cve for new keyword as it is not in database
                cve_ids_no = get_cve_list(keyword)
                for id ,des in cve_ids_no:
                    #update keyword in database and display result on web
                    query = "SELECT CVE_ID,Description,Keyword FROM vulnerabilityDB WHERE CVE_ID LIKE '{}' ".format(id)
                    cur.execute(query)
                    result2 = cur.fetchall()
                    result.append(result2[0])
                    existing_keyword = result2[0][2]
                    if bool(existing_keyword):
                        if keyword not in existing_keyword:
                            new_keyword = keyword + ',' + existing_keyword
                            new_keyword = new_keyword.lower()
                            cur.execute('''UPDATE vulnerabilityDB SET Keyword = ? WHERE CVE_ID = ? ''',(new_keyword,id))
                            my_file = open("keyword.txt", "a")
                            my_file.write("\n"+new_keyword+",")
                            my_file.close()
                    else:
                        keyword.lower()
                        cur.execute(''' UPDATE vulnerabilityDB SET Keyword = ? WHERE CVE_ID = ? ''',(keyword,id))
                    conn.commit()
            except:
                #No record found for this keyword
                my_file = open("black_listed_keys.txt", "a")
                my_file.write("\n" + keyword + ",")
                my_file.close()
                conn.close()
                return render_template('no_search_result.html')
    else:
        return render_template('no_search_result.html')
    result = sorted(result , reverse = True)
    conn.close()
    return render_template('search_result.html', searchresult=result)

@app.route('/cve/<string:cve_id_no>')
def cve_details(cve_id_no):
    link = "https://www.cvedetails.com/cve/{}/".format(cve_id_no)
    page = requests.get(link)
    soup = BeautifulSoup(page.content, 'html.parser')
    a = soup.find_all('div', class_='cvedetailssummary')
    data = a[0].text.split('\n')
    des = data[1][1:-1]
    p_date = data[3][16:26]
    u_date = data[3][46:-1]
    b = soup.find_all('table', id='cvssscorestable')
    cvss_score = b[0].find_all('tr')[0].text.split('\n')[-2]
    cwe_id = b[0].find_all('tr')[-1].text.split('\n')[-2]

    cvedata = [cve_id_no, des, p_date, u_date, cvss_score, cwe_id]

    return render_template("cve_id_details.html" , cvedetails = cvedata)

@app.route('/admin/signup',methods = ['GET','POST'])
def admin_signup():
    if request.method == 'POST':
        form = request.form
        username = form['username']
        email = form['email']
        password = form['password']
        confirm_password = form['confirm_password']

        if password == confirm_password:
            salt = os.urandom(32)
            key = hashlib.pbkdf2_hmac('sha512', password.encode(), salt, 10000)
            storage = salt + key
            conn = sqlite3.connect('vulnerabilityDB.sqlite')
            cur = conn.cursor()
            cur.execute('''INSERT INTO adminhash(username,email,hash)
                        VALUES (?,?,?)''', (username, email, storage))
            conn.commit()
            conn.close()
            return redirect(url_for("admin_login"))
    return render_template('adminsignup.html')

@app.route('/admin/login/',methods = ['GET','POST'])
def admin_login():
    if request.method == 'POST':
        session.pop('user', None)
        form = request.form
        username = form['username']
        password = form['password']
        try:

            conn = sqlite3.connect('vulnerabilityDB.sqlite')
            cur = conn.cursor()

            query = "SELECT hash FROM adminhash WHERE username LIKE '{}' ".format(username)
            cur.execute(query)
            result = cur.fetchall()
            conn.close()
            salt = result[0][0][:32]
            key = hashlib.pbkdf2_hmac('sha512', password.encode(), salt, 10000)

            if key == result[0][0][32:]:
                session['user'] = request.form['username']
                return redirect(url_for("admin"))
            else:
                return render_template('adminlogin.html')
        except:
            pass
    else:
        pass
    return render_template('adminlogin.html')

@app.route('/admin/')
def admin():
    if g.user:
        conn = sqlite3.connect('vulnerabilityDB.sqlite')
        cur = conn.cursor()
        cur.execute('''SELECT COUNT(*) FROM vulnerabilityDB WHERE Keyword IS NOT NULL''')
        labeled = cur.fetchone()[0]
        cur.execute('''SELECT COUNT(*) FROM vulnerabilityDB ''')
        total = cur.fetchone()[0]
        not_labeled = int(total) - int(labeled)
        conn.close()
        return render_template("admin.html" , total = total,labeled = labeled ,not_labeled = not_labeled)
    return redirect(url_for('admin_login'))

@app.before_request
def before_request():
    g.user = None
    if 'user' in session:
        g.user = session['user']

@app.route('/dropsession/')
def dropsession():
    session.pop('user', None)
    return redirect(url_for('admin_login'))

@app.route('/updatedb/')
def update_db():
    if g.user:
        #custom_update()
        return render_template("admin.html")
    return redirect(url_for('admin_login'))

@app.route('/changepassword/',methods = ['GET','POST'])
def change_password():
    if g.user:
        try:
            if request.method == 'POST':
                form = request.form
                username = form['username']
                current_password = form['current_password']
                new_password = form['new_password']
                cfm_new_password = form['confirm_new_password']
                conn = sqlite3.connect('vulnerabilityDB.sqlite')
                cur = conn.cursor()
                query = "SELECT hash FROM adminhash WHERE username LIKE '{}' ".format(username)
                cur.execute(query)
                result = cur.fetchall()
                salt = result[0][0][:32]
                key = hashlib.pbkdf2_hmac('sha512', current_password.encode(), salt, 10000)
                if key == result[0][0][32:]:
                    if new_password == cfm_new_password:
                        salt = os.urandom(32)
                        key = hashlib.pbkdf2_hmac('sha512', new_password.encode(), salt, 10000)
                        storage = salt + key
                        cur.execute('''UPDATE adminhash SET hash = ? WHERE username = ? ''',(storage,username))
                        conn.commit()
                        conn.close()
                        return redirect(url_for("admin_login"))
                conn.close()
        except:
            pass
    else:
        return redirect(url_for('admin_login'))
    return render_template("changepassword.html")

@app.route("/recoverpassword/", methods = ['GET','POST'])
def rec_password():
    if request.method == 'POST':
        form = request.form
        username = form['username']
        email = form['email']
        conn = sqlite3.connect('vulnerabilityDB.sqlite')
        cur = conn.cursor()
        print(username)
        cur.execute('''SELECT email FROM adminhash WHERE username = ? ''', (username,))
        result = cur.fetchone()
        conn.close()
        if result[0] == email:
            send_recovery_email(username,email)
            return redirect(url_for('admin_login'))

    return render_template("reset_password.html")

@app.route('/contact', methods = ['GET','POST'])
def contact():
    return render_template('contact.html')

if __name__ == "__main__":
    app.run(debug = True)