import codecs
import requests
import re
from bs4 import BeautifulSoup
import os
import pandas 
import xml.etree.ElementTree as ET
from lib.utils import get_ua
from lib.sqlite3_operate import cve_exists,cve_exists_where,create_table

def rss_xml():
    """监控exploit-db.com的rss数据
    :return rss_add:tuple_dict, exploit added in exploit-db.com, e.g. {eid:(title,link,des,date,guid)}
    """
    r=requests.get('https://www.exploit-db.com/rss.xml')
    rss=r.content
    root=ET.fromstring(rss)
    rss_add=dict()
    for channel in root.findall('channel'):
        for item in channel.findall('item'):
            title=item.find('title').text
            link=item.find('link').text
            des=item.find('description').text
            date=item.find('pubDate').text
            guid=item.find('guid').text
            ri=re.search(r'(\d+)',guid)
            if ri:
                eid=ri.group(0)
                rss_add[eid]=(title,link,des,date,guid)
    return rss_add

def expdb_parser(eid):
    """
    单次解析，返回单条数据
    :param eid:
    :return exp_values:dict, e.g. {'EDB-ID:': '37074', 'CVE:': '', 'Author:': 'Panagiotis Vagenas', 'Type:': 'webapps', 'Platform:': 'PHP', 'Date:': '2015-05-21'}
    :return cve_values:dict, e.g. {'EDB-ID:': '37074', 'CVE:': '2015-4039;2015-4038', 'Author:': 'Panagiotis Vagenas', 'Type:': 'webapps', 'Platform:': 'PHP', 'Date:': '2015-05-21'}
    """
    exp_values=dict()
    cve_values=dict()
    url='https://www.exploit-db.com/exploits/'+str(eid)
    r=requests.get(url,headers={'User-Agent':get_ua()},timeout=10)
    html=r.content
    if r.status_code==200 and b'404 Page Not Found' not in html:
        soup=BeautifulSoup(html,'html.parser')
        for div in soup.find_all('div',class_='col-6 text-center'):
            exp_value=list(div.stripped_strings)
            if len(exp_value)>=2:
                exp_values[exp_value[0]]=";".join(exp_value[1:])
            else:
                exp_values[exp_value[0]]=''
        if re.search(r'(\d)+',exp_values['CVE:']):
            cve_values=exp_values
        return exp_values,cve_values
    else:
        print("[!] WARNING:exploit-db response status code is %s" %r.status_code)
        return False,False

def offensive_download():
    try:
        url='https://raw.githubusercontent.com/offensive-security/exploitdb/master/files_exploits.csv'
        r=requests.get(url)
        with codecs.open('data/exploit-db/offensive_expfile.csv','wb') as f:
            f.write(r.content)
        return True
    except Exception as e:
        print("[Error] DOWNLOAD %s error,please modified hosts to access to raw.githubusercontent.com!" %url)
        return False

def expdb_exists():
    """
    从offensive Security github file exploits取最新的全量edb-id
    """
    if os.path.exists('data/exploit-db/offensive_expfile.csv'):
        exps=pandas.read_csv('data/exploit-db/offensive_expfile.csv',engine='python',error_bad_lines=False)
        return list(exps['id'].values)
    else:
        return False

def expdb_exp_add(reparse=False):
    """
    从exploit-db.com中提取CVE exp label，弥补不足
    更新策略：增量覆盖
    :return cves:tuple_dict, exploit-db.com exp data with cve number, e.g. {'48587': ('48587', '2020-12712', 'Sander Ubink', 'remote', 'Multiple', '2020-06-15'), '48590': ('48590', '2020-5515', 'BillyV4', 'webapps', 'PHP', '2020-06-16')}
    """
    print('[+] Searching exp added from remote exploit-db')
    create_table(db='exp2',table='expdb',key=['edb_id','cve_id','author','type','platform','date'],primary_key='edb_id')
    so,exist_eid=cve_exists(db='exp2',table='expdb',key=['edb_id'])
    if reparse:
        all_eid=expdb_exists()
        all_eid=[str(i) for i in all_eid]
        add_eid=list(set(all_eid).difference(set(exist_eid)))
    else:
        rss_add=rss_xml()
        rss_eid=list(rss_add.keys())
        add_eid=list(set(rss_eid).difference(set(exist_eid)))

    print("[+] Add eid:%s" %add_eid)
    
    exps=dict()
    cves=dict()
    for eid in add_eid:
        try:
            exp_value,cve_value=expdb_parser(eid)
            if exp_value:
                exps[eid]=tuple(exp_value.values())
            else:
                pass
            if cve_value:
                cves[eid]=tuple(cve_value.values())
            else:
                pass
        except Exception as e:
            print("[!] DOWNLOAD ERROR %s error:%s" %(eid,repr(e)))
    
    if exps:
        sql='replace INTO expdb (edb_id, cve_id, author, type, platform, date) VALUES (?, ?, ?, ?, ?, ?)'
        so.executemany(sql,exps.values())
        so.close()
    else:
        print("[!] exploit-db updated not found")
    
    return cves

def expdb_exp_all(): 
    """
    从本地数据库中提取explo-db.com所有有exp标记的cve
    :return exp_all:list, the whole cve id list with exp label from exploit-db.com, e.g. ['2003-0132', '2003-0132', '2003-0276']
    """
    print('[+] Searching all exp from local exploit-db database')
    exp_all=[]
    so,exist_cid=cve_exists_where(db='exp2',table='expdb',key=['cve_id'],where='cve_id not in ("N/A","")')
    for cid in exist_cid:
        if ';' in cid:
            exp_all.extend(cid.split(';'))
        else:
            exp_all.append(cid)
    return exp_all

if __name__=='__main__':
    pass
