#!/usr/bin/env python
# coding : utf-8
import argparse
import os.path
from pathlib import Path
import sys
import logging
import queue
import threading
import time
import redis


from PNL.worker import *


def percentage(part, whole):
    return 100 * float(part) / float(whole)


def main():
    g_emailonly_credential = 0
    g_noemail_credential = 0
    g_total_credential = 0
    g_total_line = 0

    parser = argparse.ArgumentParser(description='Leak parser written in Python')
    parser.add_argument('-c', '--channel', required=True, help='Redis channel to publish to')
    parser.add_argument('-H', '--host', default='127.0.0.1', help='Redis host to connect to')
    parser.add_argument('-P', '--port', default=6379, type=int, help='Redis port to connect to')
    parser.add_argument('-D', '--debug', action='store_true', help='Launch the command in debug mode')
    parser.add_argument('-d', '--directory', required=True, help='Directory to analyze')
    parser.add_argument('-w', '--worker', required=False, default=1, type=int, help='Number of workers to create')
    parser.add_argument('-e', '--extensions', nargs='+', choices=['noext', 'txt', 'zip'], required=True,
                        help='File extensions to parse')
    args = parser.parse_args()

    if not os.path.exists(args.directory):
        print('Supplied directory doesn\'t exist', file=sys.stderr)
        exit(1)

    log_file = os.path.join(args.directory, 'pnl.log')
    my_logging_format = '%(asctime)s;%(levelname)s;%(threadName)s;%(message)s'

    if args.debug:
        logging.basicConfig(filename=log_file, level=logging.DEBUG, format=my_logging_format)
    else:
        logging.basicConfig(filename=log_file, level=logging.INFO, format=my_logging_format)

    start = time.time()
    logger = logging.getLogger('pnl')
    fifo_in = queue.Queue()
    fifo_out = queue.Queue()
    fifo_stats = queue.Queue()
    files_event = threading.Event()
    cred_event = threading.Event()
    stop_event = threading.Event()
    condition_credential = threading.Condition()
    condition_redis = threading.Condition()
    # This is the default number of connection in redis pool
    redis_thread_number = 4

    # Connect to Redis
    try:
        r = redis.StrictRedis(host=args.host, port=args.port)
    except redis.ConnectionError:
        print('Redis connection error', file=sys.stderr)
        exit(1)

    try:
        # Handle conditions
        condition_credential.acquire()
        condition_redis.acquire()

        logger.info('Creating %d CredentialWorker' % args.worker)
        for worker_id in range(args.worker):
            t = CredentialWorker(worker_id, stop_event, files_event, condition_credential, fifo_in, fifo_out,
                                 fifo_stats)
            t.start()
            logger.debug('Created WorkerThread %d', t.ident)

        print('Enumerating files')
        logger.info('Enumerating files')
        for o in Path(args.directory).glob('**/*'):
            if o.is_file():
                if o.suffix == '' and 'noext' in args.extensions:
                    fifo_in.put(str(o))
                elif o.suffix in args.extensions:
                    fifo_in.put(str(o))
        logger.info('Found %d files to analyze' % fifo_in.qsize())

        # Indicate to CredentialWorker that the filesystem has been analyzed
        print('Filesystem analysis is done')
        files_event.set()

        # Check if there is a subscriber
        while r.pubsub_numsub(args.channel)[0][1] == 0:
            print('Please, launch a subscriber, then press enter to continue')
            input('...')

        print('Credential analysis started')

        logger.info('Creating %d RedisWorker' % redis_thread_number)
        for worker_id in range(redis_thread_number):
            t = RedisWorker(worker_id, stop_event, cred_event, condition_redis, r, fifo_out, args.channel)
            t.start()
            logger.debug('Created RedisThread %d', t.ident)

        logger.info('Waiting for CredentialWorker to finish')
        waiting_workers = args.worker
        while waiting_workers > 0:
            condition_credential.wait()
            try:
                a, b, c, d = fifo_stats.get()
            except queue.Empty:
                raise PNLNoStatsReceived
            g_total_line += a
            g_total_credential += b
            g_noemail_credential += c
            g_emailonly_credential += d
            waiting_workers -= 1
        condition_credential.release()

        # Indicate to RedisWorker that all credentials have been extracted
        print('Credential analysis is done')
        cred_event.set()

        logger.info('Waiting for RedisWorker to finish')
        waiting_workers = redis_thread_number
        while waiting_workers > 0:
            condition_redis.wait()
            waiting_workers -= 1
        condition_redis.release()

        print('All credentials have been sent to Redis')

    except KeyboardInterrupt:
        logger.info('PNL has been interrupted manually. Workers will be stopped')
        stop_event.set()
        try:
            condition_credential.release()
            condition_redis.release()
        except RuntimeError:
            pass

    # Display statistics
    if g_total_line:
        g_total_credential_percent = percentage(g_total_credential, g_total_line)
        g_noemail_credential_percent = percentage(g_noemail_credential, g_total_line)
        g_emailonly_credential_percent = percentage(g_emailonly_credential, g_total_line)
        print('Total lines analyzed:', g_total_line)
        print('Credentials found:', g_total_credential, '({:.2f})'.format(g_total_credential_percent))
        print('Credentials with email only:', g_emailonly_credential,
              '({:.2f})'.format(g_emailonly_credential_percent))
        print('Credentials with no email:', g_noemail_credential, '({:.2f})'.format(g_noemail_credential_percent))
        logger.info('Line:%d Creds:%d (%.2f) NoEmail:%d (%.2f) EmailOnly:%d (%.2f)', g_total_line, g_total_credential,
                    g_total_credential_percent, g_noemail_credential, g_noemail_credential_percent, g_emailonly_credential,
                    g_emailonly_credential_percent)

    # Execution time
    end = time.time()
    hours, rem = divmod(end - start, 3600)
    minutes, seconds = divmod(rem, 60)
    print('')
    print('Execution time:', '{:0>2}:{:0>2}:{:05.2f}'.format(int(hours), int(minutes), seconds))
    logger.info('Execution time:{:0>2}:{:0>2}:{:05.2f}'.format(int(hours), int(minutes), seconds))

    logger.info('Workers have finished')


if __name__ == "__main__":
    main()
