## -*- coding: utf-8 -*-
##
## ditg-control-server.py
##
## Author:   Toke Høiland-Jørgensen (toke@toke.dk)
## Date:     28 marts 2014
## Copyright (c) 2014, Toke Høiland-Jørgensen
##
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program.  If not, see <http://www.gnu.org/licenses/>.

import sys, os, optparse, subprocess, time, tempfile, shutil, json
import random, string, signal, traceback, time, hmac, hashlib

from datetime import datetime

try:
    from defusedxml.xmlrpc import monkey_patch
    monkey_patch()
    del monkey_patch
    XML_DEFUSED=True
except ImportError:
    XML_DEFUSED=False

try:
    # python 2
    from SimpleXMLRPCServer import SimpleXMLRPCServer
except ImportError:
    from xmlrpc.server import SimpleXMLRPCServer

# Nicked from the subprocess module -- for closing open file descriptors
try:
    MAXFD = os.sysconf("SC_OPEN_MAX")
except:
    MAXFD = 256

ALPHABET = list(string.ascii_lowercase) + list(string.digits)

class AlarmException(Exception):
    pass


parser = optparse.OptionParser(description="Control server for D-ITG receive component session negotiation.")

parser.add_option('--insecure-xml', action='store_true', dest='INSECURE_XML', default=False,
                  help="Run even though the defusedxml module is unavailable. WARNING: Will leave "
                  "the server open to entity expansion attacks!")
parser.add_option('-a', '--address', action='store', type='string', dest='BIND_ADDRESS',
                  default='localhost', help="Address to bind to. Default: localhost.")
parser.add_option('-p', '--port', action='store', type='int', dest='BIND_PORT', default='8000',
                  help="Bind port. Default: 8000.")
parser.add_option('-A', '--itg-address', action='store', type='string', dest='ITG_ADDRESS',
                  default=None, help="Address to bind ITGRecv to. Default: Same as --address.")
parser.add_option('-s', '--start-port', action='store', type='int', dest='START_PORT',
                  default=9000, help="Start port for ITGRecv control socket binds (default 9000).")
parser.add_option('-t', '--max-test-time', action='store', type='int', dest='MAX_TEST_TIME',
                  default=7200, help="Maximum test time allowed. Default 7200 seconds (two hours).")
parser.add_option('-m', '--max-instances', action='store', type='int', dest='MAX_INSTANCES',
                  default=100, help="Maximum number of running instances before new requests are denied (default 100).")
parser.add_option('-S', '--secret', action='store', type='string', dest='SECRET',
                  default="", help="Secret for request authentication. Default: ''.")
parser.add_option('--debug', action='store_true', dest='DEBUG',
                  default=False,  help="Debug mode: Don't delete temporary files.")

class DITGManager(object):
    datafile_pattern = "%s.json"

    def __init__(self, bind_address, start_port, max_test_time, max_instances, secret,
                 debug=False):
        self.working_dir = tempfile.mkdtemp(prefix='ditgman-')
        self.seen = {}
        self.children = []
        self.toplevel = True
        self.bind_address = bind_address
        self.max_test_time = max_test_time
        self.max_instances = max_instances
        self.hmac = hmac.new(secret.encode(), digestmod=hashlib.sha256)
        self.start_port = self.current_port = start_port
        self.id_length = 20
        self.debug = debug

        signal.signal(signal.SIGINT, self._exit)
        signal.signal(signal.SIGTERM, self._exit)
        signal.signal(signal.SIGALRM, self._alarm)
        signal.signal(signal.SIGCHLD, self._collect_garbage)

    def get_test_results(self, test_id):
        """Get the results of a previously initiated test. The sole parameter is
        the test_id returned from request_new_test().

        The return parameter is a dictionary with the following entries:

        status: 'OK' if everything went well, 'Error' otherwise.

        message: Set if status is 'Error'; contains an error message.

        data: The output of ITGDec -c with the interval requested at initiation.

        utc_offset: UTC timestamp which the time values in the return data is offset by.
        """
        self._collect_garbage()
        test_id = str(test_id)
        if len(test_id) != self.id_length:
            return {'status': 'Error', 'message': "Invalid test id: '%s'." % test_id}
        for c in test_id:
            if not c in ALPHABET:
                return {'status': 'Error', 'message': "Invalid test id: '%s'." % test_id}
        filename = os.path.join(self.working_dir,  self.datafile_pattern % test_id)
        if not os.path.exists(filename):
            return {'status': 'Error', 'message': "Data for test ID '%s' not found." % test_id}
        with open(filename, 'rt') as fp:
            data = json.load(fp)
            return data

    def request_new_test(self, duration, interval, hmac_hex, raw_data = False):
        """Request a new test instance.

        This will allocate a new ITGRecv instance and parse the log file it produces.
        The results of this can then be retrieved with get_test_results() after the
        test has run.

        The parameters are:

        duration: Requested test duration in seconds. The ITGRecv instance will be killed
                  after this time has passed (+ a grace period of five seconds).

        interval: The requested interval for data points, in milliseconds (passed to ITGDec).

        hmac_hex: A hexadecimal HMAC-SHA256 of the two other parameters computed by
                  concatenating their ASCII representations. The HMAC secret is configured
                  by the operator of the control server instance.

        raw_data: Whether to store and return the raw text log from ITGDec (i.e. the output of
                  ITGDec -l)

        The return value is a dictionary with the following keys:

        status: 'OK' if everything went well, 'Error' otherwise.

        message: Set if status is 'Error'; contains an error message.

        test_id: The assigned test ID, to be passed to get_test_results() after the duration
                 has expired.

        port:    The control server port of the ITGRecv instance. The sender is expected to
                 use port+1 for the data connection.
        """
        self._collect_garbage()
        duration = int(duration)
        interval = int(interval)
        hmac = self.hmac.copy()
        hmac.update(str(duration).encode())
        hmac.update(str(interval).encode())
        if hmac.hexdigest() != hmac_hex:
            return {'status': 'Error', 'message': "HMAC authentication failure."}
        if duration <= 0 or interval <= 0:
            return {'status': 'Error', 'message': "Duration and interval must be positive integers."}
        if duration > self.max_test_time:
            return {'status': 'Error', 'message': "Maximum test time of %d seconds exceeded." % self.max_test_time}
        if interval > duration*1000:
            return {'status': 'Error', 'message': "Interval must be <= duration."}
        if len(self.children) >= self.max_instances:
            return {'status': 'Error', 'message': "Too many concurrent instances running. Try again later."}

        test_id = "".join(random.sample(ALPHABET, self.id_length))
        # Need one port for control, one for data (if the data stream is TCP).
        port = self.current_port
        self.current_port += 2
        return self._spawn_receiver(test_id, duration, interval, port, raw_data)


    def _clean_fork(self, output=None):
        pipe_r, pipe_w = os.pipe()
        pid = os.fork()
        if pid:
            self.children.append(pid)
            os.close(pipe_w)
            return pipe_r, False
        else:
            try:
                self.children = []
                self.toplevel = False
                os.chdir(self.working_dir)
                sys.stdin.close()
                # Move the pipe fd to position 3, close everything above it
                os.dup2(pipe_w, 3)
                os.closerange(4, MAXFD)
                pipe_w = 3

                if output is not None:
                    with open(os.path.join(self.working_dir, output), "w") as fp:
                        os.dup2(fp.fileno(), 1)
                        os.dup2(fp.fileno(), 2)
                else:
                    sys.stdout.close()
                    sys.stderr.close()
            except:
                try:
                    traceback.print_exc()
                finally:
                    os._exit(1)
            return pipe_w, True

    def _unlink(self, filename):
        if self.debug:
            return
        try:
            os.unlink(filename)
        except:
            pass

    def _spawn_receiver(self, test_id, duration, interval, port, raw_data):
        stdout = "%s.recv.out" % test_id
        datafile = self.datafile_pattern % test_id
        error = False
        pipe, child = self._clean_fork(stdout)
        if child:
            try:
                ret = self._run_receiver(test_id, duration, interval, port, pipe, raw_data)
            except Exception as e:
                traceback.print_exc()
                ret = {'status': 'Error', 'message': str(e)}
                error = True
            try:
                # Write data to temp file, atomically rename to final file name
                with open('%s.tmp' % datafile, 'wt') as fp:
                    json.dump(ret, fp)
                    os.rename('%s.tmp' % datafile, datafile)
            finally:
                if not error:
                    self._unlink(stdout)
                    os._exit(0)
                os._exit(1)
        else:
            try:
                val = os.read(pipe, 1024).decode()
                if val == 'OK':
                    return {'status': 'OK', 'test_id': test_id, 'port': port}
                else:
                    return {'status': 'Error', 'message': val}
            finally:
                os.close(pipe)



    def _run_receiver(self, test_id, duration, interval, port, pipe, raw_data):
        logfile = '%s.log' % test_id
        txtlog = '%s.log.txt' % test_id
        outfile = '%s.dat' % test_id
        ret = {'status': 'OK'}

        # Run ITGRecv; does not terminate after the end of the test, so wait for
        # the test duration + a grace time of 5 seconds, then kill the process
        try:
            proc = subprocess.Popen(['ITGRecv',
                                     '-l', logfile,
                                     '-I',
                                     '-a', self.bind_address,
                                     '-Sp', str(port)])

        # Signal back to the parent whether or not we successfully spawned the receiver.
        except OSError as e:
            os.write(pipe, str(e).encode())
        else:
            time.sleep(0.1)
            try:
                w = os.waitpid(proc.pid, os.WNOHANG)
            except OSError as e:
                w = (0,str(e))
            if w != (0,0):
                ret = {'status': 'Error', 'message': 'ITGRecv exited immediately with code %s.' % w[1]}
                os.write(pipe, ret['message'].encode())
            else:
                self.children.append(proc.pid)
                os.write(pipe, b'OK')

        os.close(pipe)

        # Use an alarm signal for the timeout; means we don't have to wait for the
        # entire test duration if the ITGRecv process terminates before then
        signal.alarm(duration + 5)
        try:
            retval = proc.wait()
            if retval > 0:
                raise Exception("ITGRecv non-zero exit")
        except AlarmException:
            proc.terminate()
        signal.alarm(0)

        # Call ITGDec on the log file output, read in the resulting data
        if not os.path.exists(logfile):
            raise Exception('No data file produced')

        try:
            subprocess.check_call(['ITGDec', logfile,
                                   '-c', str(interval), outfile,
                                   '-l', txtlog])
            with open(outfile, 'rt') as fp:
                ret['data'] = fp.read()

            if not ret['data']:
                # By not raising an exception, the stdout output file is still
                # cleared, but an error is reported to the caller. This is to
                # prevent output files from sticking around for, e.g., users
                # that request a test, but then cannot get packets through; the
                # output is not useful to diagnose that anyway.
                ret = {'status': 'Error', 'message': 'Empty data set.'}

            # Read start of text log file to get timestamp of first received packet
            with open(txtlog, 'rt') as fp:
                if raw_data:
                    ret['raw'] = fp.read()
                    data = ret['raw'][:1024]
                else:
                    data = fp.read(1024)
                try:
                    idx_s = data.index('rxTime>') + len('rxTime>')
                    idx_e = data.index('Size>')
                    t,microsec = data[idx_s:idx_e].split(".")
                    h,m,s = t.split(":")
                    dt = datetime.utcnow().replace(hour=int(h),
                                                   minute=int(m),
                                                   second=int(s),
                                                   microsecond=int(microsec))
                    ret['utc_offset'] = float(time.mktime(dt.timetuple())) + dt.microsecond / 10**6
                except Exception as e:
                    ret['utc_offset'] = None

            return ret
        finally:
            # Clean up temporary files and exit
            for f in logfile, txtlog, outfile:
                self._unlink(f)


    def _collect_garbage(self, *args):
        for p in self.children:
            try:
                if os.waitpid(p, os.WNOHANG) != (0,0):
                    self.children.remove(p)
            except OSError:
                self.children.remove(p)
        # When no more children are alive, reset the current port
        if not self.children:
            self.current_port = self.start_port
        if not self.toplevel:
            return

        try:
            files = os.listdir(self.working_dir)
            for f in files:
                if not f.endswith('.json'):
                    continue
                p = os.path.join(self.working_dir, f)
                s = os.stat(p)
                td = datetime.now() - datetime.fromtimestamp(s.st_mtime)
                if td.seconds > 300:
                    self._unlink(p)
        except:
            traceback.print_exc()

    def _alarm(self, signum, frame):
        raise AlarmException()

    def _exit(self, signum, frame):
        for p in self.children:
            try:
                os.kill(p, signum)
            except OSError:
                pass

        # Do not call normal exit handlers when in a subprocess
        if not self.toplevel:
            os._exit(0)

        # SIGTERM is successful exit, others are not
        if signum == signal.SIGTERM:
            sys.exit(0)
        sys.exit(1)

    def __del__(self):
        if self.toplevel and not self.debug:
            shutil.rmtree(self.working_dir, ignore_errors=True)


def run():
    options,args = parser.parse_args()
    if not XML_DEFUSED and not options.INSECURE_XML:
        sys.stderr.write("XML EXPANSION ATTACK VULNERABILITY DETECTED. ABORTING!\n"
                         "Run with --insecure-xml to run anyway (will leave the server vulnerable!)\n")
        sys.exit(1)

    server = SimpleXMLRPCServer((options.BIND_ADDRESS, options.BIND_PORT), allow_none=True)
    manager = DITGManager(bind_address = options.ITG_ADDRESS or options.BIND_ADDRESS,
                          start_port = options.START_PORT,
                          max_test_time = options.MAX_TEST_TIME,
                          max_instances = options.MAX_INSTANCES,
                          secret = options.SECRET,
                          debug = options.DEBUG,)
    server.register_instance(manager)
    server.register_introspection_functions()
    server.serve_forever()



if __name__ == "__main__":
    run()
