#!/usr/bin/env python3

import time
import argparse
import signal
import logging

from leaky_diode import LeakyServer

# Max allowed concurrent requests
MAX_OPEN_CONNECTIONS = 10


class ExitHandler:

    def __init__(self, leaky_server):
        self.server = leaky_server
        signal.signal(signal.SIGTERM, self.exit)
        signal.signal(signal.SIGINT, self.exit)

    def exit(self, signum, frame):
        self.server.stop()
        exit(0)


if __name__ == '__main__':
 
    description = """Leaky Diode is a data exfiltration test tool for data diodes"""
    parser = argparse.ArgumentParser(description=description)
    
    parser.add_argument('host', metavar='host', type=str, help='Remore host address')
    parser.add_argument('port', metavar='port', type=int, help='Remote host port')
    parser.add_argument('secret', metavar='secret_string', 
        default='a secret string',
        help="Attack mode 'flow' or 'close' (default: %(default)s)"
    )    
    parser.add_argument('--verbose', '-v',
	dest='verbose',
	action='store_true',
	help='Show debugging messages'
    )

    args = parser.parse_args()
    
    # Change logging level in verbose mode.
    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)
    
    # Listening address and port
    server = LeakyServer(
        args.host, args.port, 
        bytes(args.secret, encoding='utf-8'),
        max_connections=MAX_OPEN_CONNECTIONS
    )
    server.start()
    exit_handler = ExitHandler(server)

    # Wait forever
    while True:
        time.sleep(10)

