#!/usr/bin/env python3

import argparse
import validators
import requests
from bs4 import BeautifulSoup, SoupStrainer
import re
import time
import os
from datetime import datetime
import signal
import sys
import urllib3
from base64 import b64encode
from multiprocessing import Queue, Process, current_process, Manager

from requests.packages.urllib3.exceptions import InsecureRequestWarning

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
requests.packages.urllib3.disable_warnings(InsecureRequestWarning)

class bcolors:
    HEADER = '\033[95m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'
    ENDC = '\033[0m'

# shared variables
manager = Manager()
results = manager.dict()
queue = manager.Value('i', 0)
checked = manager.Value('i', 0)
status_r = manager.dict()
headers = manager.dict()
q_in = manager.Queue()

def save_report(info=False, summary=False):

    if info:
        print("INFO: all results: {0}, checked: {1}, queue: {2}" .format(str(len(results)), str(checked.value),
                                                                         str(queue.value)))

    # make directories
    report_txt = workdir + "report-%d.txt" % os.getpid()
    if not os.path.exists(workdir):
        os.makedirs(workdir)

    # save report
    print(report_txt)
    f = open(report_txt,'w')
    f.write("Performance report:\n")
    f.write("-------------------\n\n")
    f.write("Site: " + base_url + "\n\n")
    f.write("All pages: " + str(len(results.copy())) + "\n")
    f.write("Checked: " + str(checked.value) + "\n")
    f.write("Queue: " + str(queue.value) + "\n")
    f.write("Http codes:\n")

    # all codes
    for key in sorted(status_r):
        f.write( "- " + str(key) + ": " + str(status_r[key]) + "\n")

    # detailed wrong codes
    header = 0
    for key in status_r:
        if int(key) >= 400:

            if header == 0:
                f.write( "\nWrong http codes:\n")
                header = 1

            # show link
            for k in results:
                if results[k]['code']:
                    if int(results[k]['code']) == key:
                        f.write("- " + str(k) + " " + str(results[k]['code']) + "\n")

    # slow responses
    header = 0
    for k in results:
        if results[k]['time'] >= 1 and results[k]['time'] < 2:
            if header == 0:
                f.write( "\nSlow responses (1-2s):\n")
                header = 1
            f.write("- " + str(k) + " " + str(results[k]['time']) + "s\n")

    # critical responses
    header = 0
    for k in results:
        if results[k]['time'] > 2:
            if header == 0:
                f.write( "\nVery slow responses (more than 2s):\n")
                header = 1
            f.write("- " + str(k) + " " + str(results[k]['time']) + "s\n")
    f.flush()
    f.close()

    if summary:
        file = open(report_txt, "r")
        print(file.read())
        print(report_txt)

    if debug:
        print("DEBUG: results: {0} , queue: {1}" .format(str(len(results)), str(sys.getsizeof(results))))

def signal_handler(signal, frame):
    if current_process().name == 'MainProcess':
        save_report(summary=True)
    sys.exit(0)

def process2(q_in):

    if debug:
        print('Starting process: %s, pid: %d' % (current_process().name, current_process().pid))

    while True:
        try:
            # get from queue
            url = q_in.get(timeout=3)
        except:
            # if processes > queue
            break

        if checked.value == max_requests:
            break

        queue.value = queue.value - 1
        checked.value = checked.value + 1

        stime = time.time()
        if validators.url(url) and results[url]['code'] == 0:
            start_rtime = time.time()

        u = urllib3.util.parse_url(url)
        host = u.scheme + "://" + u.netloc
        try:
            start_time = time.time()
            s = requests.Session()
            r = s.get(url, headers=headers, verify=False, timeout=5)
            exec_time = round(time.time()-start_time,2)

        except Exception as e:
            print(bcolors.FAIL + "Error: can't fetch " + url + bcolors.ENDC)
            if 598 not in status_r:
                status_r[598] = 1
            else:
                status_r[598] = status_r[598] + 1

            results[url] = {'code': '598', 'time': 0, 'message': str(e)}
            return False

        request_time = round(time.time()-start_rtime,4)

        # set color
        if exec_time < 1:
            color = bcolors.ENDC
        elif exec_time >= 1 and exec_time < 2:
            color = bcolors.WARNING
        else:
            color = bcolors.FAIL

        # show message
        print(color + "+ {0} {1} {2}s" . format(url, r.status_code, exec_time) + bcolors.ENDC)

        start_stime = time.time()
        if r.status_code not in status_r:
            status_r[r.status_code] = 1
        else:
            status_r[r.status_code] = status_r[r.status_code] + 1
        
        html_page = r.text

        results[url] = {'time': round(time.time()-start_time, 2), 'code': r.status_code, 'header': r.headers }
        status_time = round(time.time()-start_stime, 4)

        start_soup_time = time.time()
        strainer = SoupStrainer('a')
        soup = BeautifulSoup(html_page.encode('ascii', 'ignore'), 'html.parser', parse_only=strainer)
        soup_time_2 = round(time.time()-start_soup_time, 4)

        for link_r in soup.findAll('a'):
            link = link_r.get('href')
            search = False
            start_soup_search_time = time.time()
            try:
                if re.search(r'^'+url,str(link)):
                    search = True
                elif re.search(r'^//',str(link)):
                    search = False
                elif re.search(r'^/',str(link)):
                    search = True
                    link = host + link
                else:
                    search = False
            except Exception as e:
                print(bcolors.WARNING + "Warning: " + str(e) + bcolors.ENDC)
                search = False

            soup_search_time = round(time.time()-start_soup_search_time,4)

            # add to queue
            start_soup_queue_time = time.time()
            if search is True:
                if link not in results:

                    if validators.url(link):
                        results[link] = {'code': 0, 'time': 0}
                        queue.value = queue.value + 1

                        q_in.put(link)  # dodaj link do kolejki

                    soup_queue_append_time = round(time.time()-start_soup_queue_time,4)
                    if debug:
                        print("DEBUG: link: " + str(link) + " soup_queue_append_time " +  str(soup_queue_append_time) + " " + str(search))

            soup_queue_time = round(time.time()-start_soup_queue_time,4)

        del soup
        del html_page

        soup_time = round(time.time()-start_soup_time, 4)

        if queue.value == 0:
            if debug:
                print('Stopping process: %s, pid: %d' % (current_process().name, current_process().pid))
            break


        # # save after x requests
        # if checked.value%100 == 0:
        #     self.save_report(info=True)
        #
        # function_time = round(time.time()-stime, 4)
        # if self.debug:
        #     print("DEBUG: function time: {0}" . format(function_time))
        #     print("DEBUG: request time: {0}" . format(request_time))
        #     print("DEBUG: status time: {0}" . format(status_time))
        #     print("DEBUG: soup time: {0}" . format(soup_time))
        #     print("DEBUG: soup time 2: {0}" . format(soup_time_2))
        #     print("\n")

def validate_url(url):
    if not validators.url(url):
        print(bcolors.FAIL + "Error: url is not valid." + bcolors.ENDC)
        exit(1)


if __name__ == "__main__":

    signal.signal(signal.SIGINT, signal_handler)

    parser = argparse.ArgumentParser(description='Linear Spider')
    parser.add_argument('-H', '--header', action='append', nargs='*', help='send this HTTP header '
                                                                           '(you can specify several)')
    parser.add_argument('--debug', action="store_true")
    parser.add_argument('-C', '--credentials', help='provide credentials for basic authentication (user:pass)')
    parser.add_argument('url', help='Checked site')
    parser.add_argument('-c', '--concurrency', type=int, help='Number of multiple requests to make at a time',
                        default=1)
    parser.add_argument('-n', '--requests', type=int, help='Number of requests to perform',
                        default=-1)
    args = parser.parse_args()

    debug = False

    now = datetime.today().strftime('%Y-%m-%d_%H')
    url = args.url
    base_url = args.url
    max_requests = args.requests

    # set debug
    debug = args.debug

    # set headers
    if args.header:
        for h in args.header:
            n = h[0].split(':')
            if len(n) == 2:
                k = n[0].strip()
                if n[0].lower().strip() == 'user-agent':  # overwrite if needed
                    k = 'User-Agent'
                v = n[1].strip()

                headers[k] = v
                if debug:
                    print('DEBUG: setting header: "%s: %s"' % (k, v))

    # set default user-agent
    if 'User-Agent' not in headers:
        headers['User-Agent'] = 'Mozilla/5.0 (X11; U; Linux i686; pl; rv:1.8.1.4) Gecko/20070705 ' \
                                'Firefox/2.0.0.4 (linear)'

    # set credentials
    if args.credentials:
        credentials = args.credentials.split(':')
        if len(credentials) == 2:
            user_pass = b64encode(args.credentials.encode('utf-8')).decode("ascii")
            headers['Authorization'] = 'Basic %s' % user_pass
            if debug:
                print('DEBUG: setting header "Authorization: Basic %s"' % user_pass)

    # check url
    validate_url(url)

    u = urllib3.util.parse_url(url)
    workdir = os.getenv("HOME") + "/.linear-spider/" + u.netloc + "/" + now + "/"

    ## procesy
    processes = []
    for _ in range(args.concurrency):
        p = Process(target=process2, args=(q_in,))
        p.start()
        processes.append(p)
    #
    # print("procesy utworzone")
    #
    # dodaj pierwszy link do kolejki, potem się potoczy..

    if len(results) == 0:
        results[url] = {'code': 0, 'time': 0}
        queue.value = 1
        q_in.put(url)

    for proc in processes:
        proc.join()

    time.sleep(1)

    for p in processes:
        p.terminate()

    save_report(summary=True)
