import pika
import pyautogui
import cv2
import time
import pyperclip as pyperclip
import requests
import sqlite3
import json
import cv2 as cv
import numpy
from urllib.parse import urlparse


# function for getting data from database
def data_from_database():
    global browser_name
    global os_name
    global range
    global host
    global user
    global password
    global queue
    global exchange
    global reporter_queue

    conn = sqlite3.connect('phishing.sqlite')
    c = conn.cursor()

    browser_name = c.execute('SELECT value FROM configurations WHERE key_data=\'browser_name\' LIMIT 1;').fetchone()[0]
    os_name = c.execute('SELECT value FROM configurations WHERE key_data=\'os_name\' LIMIT 1;').fetchone()[0]
    range = c.execute('SELECT value FROM configurations WHERE key_data=\'comparative_index\' LIMIT 1;').fetchone()[0]
    host = c.execute('SELECT value FROM configurations WHERE key_data=\'rabbit_mq_host\' LIMIT 1;').fetchone()[0]
    user = c.execute('SELECT value FROM configurations WHERE key_data=\'rabbit_mq_customer_user\' LIMIT 1;').fetchone()[
        0]
    password = \
    c.execute('SELECT value FROM configurations WHERE key_data=\'rabbit_mq_customer_pass\' LIMIT 1;').fetchone()[0]
    if browser_name == 'chrome':
        queue = 'phishing_new'
    else:
        queue = \
            c.execute('SELECT value FROM configurations WHERE key_data=\'rabbit_mq_customer_edge_queue\' LIMIT 1;').fetchone()[0]
    exchange = \
    c.execute('SELECT value FROM configurations WHERE key_data=\'rabbit_mq_customer_exchange\' LIMIT 1;').fetchone()[0]
    reporter_queue = \
    c.execute('SELECT value FROM configurations WHERE key_data=\'rabbit_mq_reporter_queue\' LIMIT 1;').fetchone()[0]

    conn.close()


# function for connection to RabbitMQ server
def connect_to_rabbit_mq_server():
    global credentials

    credentials = pika.PlainCredentials(user, password)
    connection = pika.BlockingConnection(pika.ConnectionParameters(
        host=host, port=5672, credentials=credentials))  # open connection
    channel = connection.channel()  # get the channel
    channel.queue_declare(queue=queue)  # set the queue
    return channel


# main function
def run_test(ch, method, properties, body):
    data = json.loads(body)
    data['system'] = 'new_system'
    data['domain'] = data['address']

    print(" [x] Starting phishing test for url: " + data['domain'])  # print to log
    ch.basic_ack(delivery_tag=method.delivery_tag)  # keep alive and send feedback
    if test_url(data, 1):  # test availability
        # if test_redirect(url_parsed):  # test redirect
        if browser_name == 'chrome':
            control_browser(os_name, data['domain'])  # control Chrome to url
        else:
            control_browser(os_name, data['address'])  # control Chrome to url
        get_image_difference(data)  # check pictures
        clear_downloads()  # clear Downloads folder


# function for deleting files in Downloads folder
def clear_downloads():
    import os
    import glob

    files = ''
    if os_name == 'macos':
        files = '/Users/' + os.getlogin() + '/Downloads/*'
    if os_name == 'windows':
        files = glob.glob('C:/Users/' + os.getlogin() + '/Downloads/*')
    for f in files:
        try:
            os.remove(f)
        except Exception as e:
            print(e)


# function for testing url availability, try 5 times and report if fails
def test_url(data, current_try):
    try:
        requests.get('http://' + data['address'], timeout=(10, 30))
        return True
    except Exception as e:
        if current_try > 5:
            print('Url ' + data['domain'] + ' cant be reached, error: ' + str(e))
            report_invalid_url(data, str(e))
            return False
        else:
            current_try += 1
            test_url(data, current_try)


# function for testing redirect | if url is different then report
def test_redirect(url):

    # test availability
    try:
        r = requests.get(url, timeout=(10, 30))

        # test redirect
        if '{uri.scheme}://{uri.netloc}'.format(uri=urlparse(url)) in r.url:
            return True
        else:
            print('Url ' + url + ' has redirection, error: ' + "Redirected url: " + r.url)
            report_invalid_redirect(url, "Redirected url: " + r.url)
            return False
    except Exception as e:
        print('Url ' + url + ' is invalid, error: ' + str(e))
        return False


# function for opening browser
def open_browser():

    if browser_name == 'chrome':

        if os_name == 'macos':
            pyautogui.hotkey('ctrl', 'shift', ']')
        if os_name == 'windows':
            pyautogui.hotkey('ctrl', 'alt', ']')

    elif browser_name == 'edge':

        if os_name == 'macos':
            pyautogui.hotkey('ctrl', 'shift', '[')
        if os_name == 'windows':
            pyautogui.hotkey('ctrl', 'alt', '[')

    time.sleep(4)


# function for controlling browser for getting screenshots
def control_browser(os_name, url):

    # open new tab
    if os_name == 'macos':
        pyautogui.hotkey('command', 't')
    if os_name == 'windows':
        pyautogui.hotkey('ctrl', 't')
    time.sleep(1)

    # type url
    if os_name == 'macos':
        pyautogui.hotkey('command', 'l')
    if os_name == 'windows':
        pyautogui.hotkey('ctrl', 'l')
    time.sleep(1)

    # fast text enter
    pyperclip.copy(str(url))
    if os_name == 'macos':
        pyautogui.hotkey('command', 'v')
    if os_name == 'windows':
        pyautogui.hotkey('ctrl', 'v')

    pyautogui.press('enter')
    time.sleep(5)

    # make screenshow
    pyautogui.screenshot('fresh_screenshot_phished_screen.png')
    time.sleep(2)

    # close tab
    if os_name == 'macos':
        pyautogui.hotkey('command', 'w')
    if os_name == 'windows':
        pyautogui.hotkey('ctrl', 'w')
    time.sleep(1)


# function for closing browser
def close_browser():
    if os_name == 'windows':
        pyautogui.hotkey('alt', 'f4')


# function for creating screenshot for saving
def generate_screenshot_for_saving(phished):
    import os

    try:
        path = ''

        if phished:
            folder = 'Phished'
        else:
            folder = 'Clean'

        if os_name == 'macos':
            path = '/Users/' + os.getlogin() + '/Desktop/Screenshots/' + folder
        elif os_name == 'windows':
            path = 'C:/Users/' + os.getlogin() + '/Desktop/Screenshots/' + folder

        if not os.path.exists(path):
            os.makedirs(path)

        pyautogui.screenshot(path + '/screenshot_' + time.strftime("%Y_%m_%d-%H_%M_%S") + '.png')
    except Exception as e:
        print(e)


# function for closing Allow / Block pop-up
def close_allow_block_popup():
    try:

        x_coordinate = 0
        y_coordinate = 0

        img_rgb = cv.imread('fresh_screenshot_phished_screen.png')
        img_gray = cv.cvtColor(img_rgb, cv.COLOR_BGR2GRAY)
        template = cv.imread('allow_block_template.png', 0)
        res = cv.matchTemplate(img_gray, template, cv.TM_CCOEFF_NORMED)
        threshold = 0.8
        loc = numpy.where(res >= threshold)
        for pt in zip(*loc[::-1]):
            x_coordinate = pt[0] + 5
            y_coordinate = pt[1] + 5

        if x_coordinate != 0 or y_coordinate != 0:
            pyautogui.moveTo(x_coordinate + 5, y_coordinate + 5)
            pyautogui.click()
            time.sleep(2)

    except Exception as e:
        print(e)


# function for defining difference percentage
def get_image_difference(data):
    image_1 = cv2.imread('fresh_screenshot_phished_screen.png', 0)
    if browser_name == 'chrome':
        image_2 = cv2.imread('template_for_phished_site3.png', 0)
    else:
        image_2 = cv2.imread('phished_template_' + browser_name + '.png', 0)

    first_image_hist = cv2.calcHist([image_1], [0], None, [256], [0, 256])
    second_image_hist = cv2.calcHist([image_2], [0], None, [256], [0, 256])

    img_hist_diff = cv2.compareHist(first_image_hist, second_image_hist, cv2.HISTCMP_BHATTACHARYYA)
    img_template_probability_match = \
        cv2.matchTemplate(first_image_hist, second_image_hist, cv2.TM_CCOEFF_NORMED)[0][0]
    img_template_diff = 1 - img_template_probability_match

    # taking only 10% of histogram diff, since it's less accurate than template method
    commutative_image_diff = (img_hist_diff / 10) + img_template_diff
    report_phishing(commutative_image_diff, data)


# function for reporting phishing
def report_phishing(comparative_index, data):
    if comparative_index < float(range):
        print("Url: " + data['domain'] + " has been phished!")
        print("Match percentage: " + str(comparative_index))
        data['error_type'] = 'phishing'
        data['error_message'] = ''
        data['browser_name'] = browser_name
        load_rabbitmq(json.dumps(data))
    else:
        close_allow_block_popup()


# function for url unavailable
def report_invalid_url(data, error):
    print("Url: " + data['domain'] + " cant be reached!")
    data['error_type'] = 'invalid'
    data['error_message'] = error
    load_rabbitmq(json.dumps(data))


# function for invalid redirect
def report_invalid_redirect(data, error):
    print("Url: " + data['domain'] + " has invalid redirect!")
    data['error_type'] = 'redirect'
    data['error_message'] = error
    load_rabbitmq(json.dumps(data))


# function for loading rabbit MQ
def load_rabbitmq(data):
    report_connection = pika.BlockingConnection(pika.ConnectionParameters(
        host=host, port=5672, credentials=credentials))  # open connection
    report_channel = report_connection.channel()  # get channel
    report_channel.queue_declare(queue=reporter_queue)  # set queue
    report_channel.basic_publish(exchange='',
                                 routing_key=reporter_queue,
                                 body=data,
                                 properties=pika.BasicProperties(delivery_mode=2))  # publish message to customers
    report_connection.close()  # close connection


# function for consuming from rabbitmq
def start_consuming():
    channel = connect_to_rabbit_mq_server()
    channel.basic_qos(prefetch_count=1)
    channel.basic_consume(run_test, queue=queue)  # run function when called
    channel.start_consuming()  # start consuming


# main function
def main():
    data_from_database()  # load data from database
    open_browser()  # open chrome

    try:
        start_consuming()
    except Exception as e:
        print(e)
        start_consuming()

    close_browser()  # close Chrome


main()  # run script