from selenium import webdriver
from selenium.webdriver.common.keys import Keys
from bs4 import BeautifulSoup
import os
import time
import re
import mysql.connector as mariadb


if __name__ == '__main__':
    driver = webdriver.Chrome('/home/kerby/projects/pyscan/chromedriver')
    mariadb_connection = mariadb.connect(user='root', password='toor', database='work')
    data=mariadb_connection.cursor()
    malware_list=["iot", "malware", "trojan", "network", "ssh", "ftp", "rootkit", "banker", "ransom", "ransomware", "ethernalblue", "exploit", "cve", "linux", "windows", "android", "mirai"]
    for malware in malware_list:
        driver.get("https://www.virustotal.com/ru/")
        driver.find_element_by_id("search-tab-chooser").click()
        driver.find_element_by_id("query").send_keys(malware)
        driver.find_element_by_id("search").click()
        time.sleep(5)
        src_code=src=BeautifulSoup(driver.page_source, 'lxml')
        all_files=[a['href'] for a in src_code.find_all('a', href=re.compile(r'^/ru/file/\w*'))]
        data.execute("SELECT link FROM malwares")
        already=[link[0] for link in data]
        for element in all_files:
            if element in already: continue;
            print("Элемент добавлен: ",(element))
            data.execute("INSERT INTO malwares (type, link) VALUES (%s, %s)",(malware, element))
        mariadb_connection.commit()
    driver.close()
    os.system("killall chromedriver")