'''
##################################################################################
###############################      INFO      ###################################

# Program that uses a csv from which it can retrieve IDs.
# Those are needed to download further information from a website. Website is called by ID. (Therefore the csv with IDs)
# The new infos will then be saved into a expanded version of the given csv.
# Copy of the old data will remain. New data will be stored in extra file.

© Maximilian Karl, 2018


##################################################################################
##################################    LOG    #####################################

# 20.04.2018 | Started work on this program.

    - Implemented new method startProgram() to start the program and handle the user-interaction
        * Therefore ask4CsvName() was neccessary.
        * Additionally a method had to be implemented. (retrieveInfoByID())

    - Within retrieveInfoByID() there now is a method called collectInfo().
        * It will return a list of infos.
        * Those infos are generated by goToWebsite() and extractInformation()  (similar to crawler)
        * extractInformation() is not doing anything currently. ---> TO DO:  how to find_all("b") & after?

    <TO DOs>
        + Add a progressbar for the fetching process.
        + Check out how to get infos out of <b> Size: </b> 65329 <BR>


# 22.04.2018 | Finished the main functions of the program.

    <TO DO's>               STATUS                   --->                       >> DONE <<
        + Add a progressbar for the fetching process.                           >> DONE <<
        + Check out how to get infos out of <b> Size: </b> 65329 <BR>           >> DONE <<

    - Implemented a progressbar similar to the one in "crawler0_5.py".

    - Found out, how to extract information from the HTML-Code on the new site.

    - Implemented a method to save the new data inside a new CSV-File.
        * User has to enter, how the file will be called (new method: askNewCsvName()).

    - Program WORKS now!

    - Cleaned up a bit and deleted unneccessary stuff.


# 24.04.2018 | Recognized an error.

    <TO DO's>
        + GIT upload!!!!
        + Fix the problem, that occurs, when the called website has not all values that are neccessary for the CSV-File.
            * The Problem is, that the website has no html-code like "SHA-1" -> will be skipped -> next record gets into false entity.

    - Fixed the Problem!
        * Added some helper-methods: getNumOfObjectsIn() & fillBy()
        * fillBy() fills all empty columns! -> No false columns anymore!


'''


import bs4 as bs
from numpy import nan
from urllib.request import Request, urlopen
import urllib.request as req
import urllib.parse as parse
import pandas as pd
import os
import sys
import csv
import math
import time


class advanced_crawler:

    path = '/Users/Maxi/Desktop/atom/python/informationSecurity/seminar/tests/advanced_crawler/databases/'
    path_crawler = '/Users/Maxi/Desktop/atom/python/informationSecurity/seminar/tests/crawler/databases/'
    url_vxVault = 'http://vxvault.net/ViriFiche.php?ID='

    def __init__(self):
        print()
        print("#####    A D V A N C E D    C R A W L E R    #####")
        print()
        self.startProgram()


    # Starts the program and handles the user interaction (and its errors).
    def startProgram(self):
        self.csv_name = self.askCsvName()
        self.end_location = self.path + self.askNewCsvName()
        # Faster (for debugging):
        # self.csv_name = "vxVault"
        # self.end_location = self.path + "vxVault_optimized2"
        self.filepath = self.path_crawler + self.csv_name
        self.df = pd.DataFrame()

        try:
            self.df = pd.read_csv(self.filepath + ".csv", index_col=None)
        except Exception:
            print("[*] ERROR! File not found!")
            self.startProgram()

        self.id_list = self.df['ID']                                                    # Returns a list of IDs from df.
        # self.id_list = [37967, 37966, 37965, 37964, 37963, 37962, 37961, 37960]       # For test purposes.
        if os.path.exists(self.end_location + ".csv"):
            self.existing_df = pd.read_csv(self.end_location + ".csv", index_col=None)
            self.id_list_existing = self.existing_df['ID']
            self.id_list = list(set(self.id_list) - set(self.id_list_existing))
            self.id_list.sort(reverse=True)
            self.new_df = self.retrieveInfoByID(self.id_list)
            self.updateData(self.new_df, self.filepath, self.id_list, self.end_location)
        else:
            self.new_df = self.retrieveInfoByID(self.id_list)                           # As init
            self.saveData(self.new_df, self.df, self.end_location)


    # Handles the user input for the wished csv-name.
    def askCsvName(self):
        self.name = str(input("[>] Enter CSV-name: "))
        if len(self.name) == 0:
            print("[*] ERROR! Enter something!")
            self.askCsvName()

        return self.name


    # Handles the user input for the wished csv-name.
    def askNewCsvName(self):
        self.name = str(input("[>] Enter name for extended CSV: "))
        if len(self.name) == 0:
            print("[*] ERROR! Enter something!")
            self.askCsvName()

        return self.name


    # Returns a dataframe with all extra-info.
    def retrieveInfoByID(self, id_list):
        self.start_time = time.time()
        self.content_size = 0
        self.new_df = []
        self.id_amount = len(id_list)
        self.counter = 0

        # Progress should be called here.
        for self.id in id_list:
            self.new_df.append(self.collectInfo(self.id))

            self.elapsed_time = self.getTime(self.start_time)
            self.content_size += 2.5
            self.counter += 1
            self.progress(self.counter, self.id_amount, " | (" + str(int(self.counter)) + ") Webpages", self.elapsed_time, self.content_size * 1024)

        return self.new_df


    # Helper method ("interface").
    def collectInfo(self, id):
        self.url = self.url_vxVault + str(id)
        self.soup = self.goToWebsite(self.url)
        self.info_list = self.extractInformation(self.id, self.soup)
        return self.info_list


    # Accesses the website and returns a BeautifulSoup - Object
    def goToWebsite(self, website):
        self.source = req.urlopen(website).read()
        self.soup = bs.BeautifulSoup(self.source, 'lxml')
        return self.soup


    # Does further work with the soup-content
    def extractInformation(self, id, soup):
        self.info_arr = []
        self.info_arr.append(id)
        self.num_objects = 0

        # Goes through every line of the webpage and searches for a specific tag (Example: <a>)
        for self.line in self.soup.get_text().split('\n'):
            self.text = self.line
            self.num_objects = self.getNumOfObjectsIn(self.info_arr)

            if 'File:' in str(self.text):
                self.info_arr.append(self.text.split(': ')[-1]) #NECCESSARY!!!
            elif 'Size:' in str(self.text):
                if self.num_objects == 2:
                    self.info_arr.append(self.text.split(': ')[-1])
                else:
                    self.info_arr = self.fillBy(self.info_arr, 2 - self.num_objects)
                    self.info_arr.append(self.text.split(': ')[-1])
            elif 'SHA-1:' in str(self.text):
                if self.num_objects == 3:
                    self.info_arr.append(self.text.split(': ')[-1])
                else:
                    self.info_arr = self.fillBy(self.info_arr, 3 - self.num_objects)
                    self.info_arr.append(self.text.split(': ')[-1])
            elif 'SHA-256:' in str(self.text):
                if self.num_objects == 4:
                    self.info_arr.append(self.text.split(': ')[-1])
                else:
                    self.info_arr = self.fillBy(self.info_arr, 4 - self.num_objects)
                    self.info_arr.append(self.text.split(': ')[-1])
            elif 'Added:' in str(self.text):
                if self.num_objects == 5:
                    self.info_arr.append(self.text.split(': ')[-1])
                else:
                    self.info_arr = self.fillBy(self.info_arr, 5 - self.num_objects)
                    self.info_arr.append(self.text.split(': ')[-1])

        return self.info_arr


    # Small helper method for extractInfo.
    def getNumOfObjectsIn(self, list):
        return len(list)


    # Fills a list with "NONE" on empty spaces (you have to tell, how many are empty).
    def fillBy(self, list, num):
        for i in range(num):
            list.append("NaN")
        return list


    # Displays the current progress of something
    def progress(self, count, total, status, elapsed_time, content_size):
        sys.stdout.flush()
        self.fill = '█'
        self.bar_len = 60
        self.filled_len = int(round(self.bar_len * count / float(total)))

        self.wps = float(count / elapsed_time)
        self.size = self.checkSizeFormat(content_size)
        self.unit = self.getUnit(content_size)
        self.speed = self.checkSizeFormat(float(content_size / elapsed_time))
        self.speed_unit = self.getUnit(float(content_size / elapsed_time)) + "/s"

        self.eta = int((total - count) / self.wps)

        # V ESTIMATED TIME OF ARRIVAL:

        self.eta_hours = 0
        self.eta_mins = 0
        self.eta_secs = 0

        # Better way of formatting! (not like down if...)
        self.eta_mins, self.eta_secs = divmod(self.eta, 60)
        self.eta_hours, self.eta_mins = divmod(self.eta_mins, 60)

        self.eta_hours = self.checkTimeFormat(self.eta_hours)
        self.eta_mins = self.checkTimeFormat(self.eta_mins)
        self.eta_secs = self.checkTimeFormat(self.eta_secs)

        # V ELAPSED TIME:

        self.hours = 0
        self.mins = 0
        self.secs = 0

        # Better way of formatting! (not like down if...)
        self.mins, self.secs = divmod(int(elapsed_time), 60)
        self.hours, self.mins = divmod(self.mins, 60)

        self.hours = self.checkTimeFormat(self.hours)
        self.mins = self.checkTimeFormat(self.mins)
        self.secs = self.checkTimeFormat(self.secs)

        self.percents = round(100.0 * count / float(total), 1)
        self.bar = str(self.fill * self.filled_len + ' ' * (self.bar_len - self.filled_len))

        output1 = "\r[*] [%s]  " % (self.bar)
        output2 = '\b%s%% %s in %sh:%sm:%ss | ETA: %sh:%sm:%ss | Samples/s: %.2f | Downloaded %.2f%s @ %d%s...      ' % (self.percents, status, self.hours, self.mins, self.secs, self.eta_hours, self.eta_mins, self.eta_secs, self.wps, self.size, self.unit, self.speed, self.speed_unit)
        sys.stdout.write(output1 + output2)
        if count == total:
            print()
            print("\n[*] DOWNLOAD successfull...")


    # Calculates the amount of the passed time (For loading bar)
    def getTime(self, start):
        self.time_passed = time.time() - start
        return self.time_passed


    # Checks the time format and automatically changes it accordingly
    def checkTimeFormat(self, time):
        if time < 10:
            time = "0" + str(time)
        else:
            time = str(time)
        return time


    # Formats the size of a file
    def checkSizeFormat(self, size):
        if size < 1024:
            return float(size)                             # BYTES
        elif (size >= 1024) & (size < 1048576):
            return float(size / 1024.0)                      # KILO-BYTES
        elif (size >= 1048576) & (size < 1073741824):
            return float(size / 1048576.0)                   # MEGA-BYTES
        else:
            return float(size / 1073741824.0)


    # Returns the unit of a certain size
    def getUnit(self, size):
        if size < 1024:
            return "B"                              # BYTES
        elif (size >= 1024) & (size < 1048576):
            return "KB"                             # KILO-BYTES
        elif (size >= 1048576) & (size < 1073741824):
            return "MB"                             # MEGA-BYTES
        else:
            return "GB"


    # Saves all the data in an extra CSV-file.
    def saveData(self, new_df, df, path):
        self.new_data = pd.DataFrame(new_df, columns=['ID1', 'NAME', 'SIZE', 'SHA-1', 'SHA-256', 'DATE'])
        df = pd.merge(df, self.new_data, left_on='ID', right_on='ID1', how='left').drop(['ID1'], axis=1)

        df.to_csv(path + '.csv', index=False)
        print()
        print("[*] SAVED Data as CSV @ ~" + path + '.csv')
        print()
        print("#####    E N D    O F    P R O G R A M   #####")
        print()


    # Updates the CSV, by adding only new mw-samples.
    def updateData(self, new_df, csv_location, list, end_location):
        self.csv_data = self.getCsvDataByID(csv_location, list)
        self.existing_data = pd.read_csv(end_location + '.csv', index_col=None)

        self.new_data = pd.DataFrame(new_df, columns=['ID1', 'NAME', 'SIZE', 'SHA-1', 'SHA-256', 'DATE'])
        self.new_df = pd.merge(self.csv_data, self.new_data, left_on='ID', right_on='ID1', how='left').drop(['ID1'], axis=1)

        print()
        self.result = pd.concat([self.new_df, self.existing_data], axis=0)
        with open(end_location + '.csv', 'w') as f:
            self.result.to_csv(f, index=False)
        print("[*] Updated CSV-File @ ~" + end_location)
        print()


    # Returns a DataFrame of the given ids in the id list within a certain location.
    def getCsvDataByID(self, location, list):
        self.df = pd.read_csv(location + '.csv')
        self.data = pd.DataFrame(columns=['ID', 'URL', 'MD5', 'IP', 'STATUS'])

        self.data = self.df.loc[0:len(list)-1]

        return self.data



advanced_crawler()
