## -*- coding: utf-8 -*-
##
## tcpdump-wrapper.py
##
## Author:   Toke Høiland-Jørgensen (toke@toke.dk)
## Date:     30 April 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 signal, sys, os, optparse, subprocess, shutil, time, socket

try:
    MAXFD = os.sysconf("SC_OPEN_MAX")
except:
    MAXFD = 256

DEFAULT_IFACE='eth1'
DEFAULT_SIZE=128
DEFAULT_DEST_DIR="/home/data"
CAPTURE_DIR="/tmp"

parser = optparse.OptionParser(description="Wrapper to start/stop tcpdump.",
                               usage="Usage: %prog [options] <start|stop> <filename>")

parser.add_option('-i', '--interface', action='store', type="string", dest='INTERFACE', default=DEFAULT_IFACE)
parser.add_option('-s', '--size', action='store', type="int", dest='SIZE', default=DEFAULT_SIZE)
parser.add_option('-d', '--dest-dir', action='store', type="string", dest='DEST_DIR', default=DEFAULT_DEST_DIR)


def start_tcpdump(filename, iface=DEFAULT_IFACE, size=DEFAULT_SIZE):
    pidfile = os.path.join(CAPTURE_DIR, "%s.pid" % filename)
    if os.path.exists(pidfile):
        sys.stderr.write("Pidfile already exists: %s.\n" % pidfile)
        sys.exit(1)
    args = ['sudo', 'tcpdump', '-n', '-i', iface, '-s', str(size), '-w', os.path.join(CAPTURE_DIR, "%s.cap" % filename)]
    pid = os.fork()
    if pid:
        with open(pidfile, "w") as fp:
            fp.write("%d\n" % pid)
        sys.exit(0)
    else:
        logfile = os.path.join(CAPTURE_DIR, "%s.log" % filename)
        fp_in = open("/dev/null", "r")
        fp_out = open(logfile, "a")
        os.dup2(fp_in.fileno(), 0)
        os.dup2(fp_out.fileno(), 1)
        os.dup2(fp_out.fileno(), 2)
        os.closerange(3, MAXFD)
        prog = args[0]
        os.execvp(prog, args)

def stop_tcpdump(filename, dest_dir=DEFAULT_DEST_DIR):
    pidfile = os.path.join(CAPTURE_DIR, "%s.pid" % filename)
    logfile = os.path.join(CAPTURE_DIR, "%s.log" % filename)
    datafile = os.path.join(CAPTURE_DIR, "%s.cap" % filename)
    dest_data = os.path.join(dest_dir, "%s.cap.gz" % filename)
    dest_log = os.path.join(dest_dir, "%s.log" % filename)
    try:
        with open(pidfile, "r") as fp:
            pid = int(fp.read())
    except (OSError, IOError, ValueError) as e:
        sys.stderr.write("Unable to read pidfile: %s.\n" % e)
        sys.exit(1)

    try:
        subprocess.check_call("sudo kill %d" % pid, shell=True)
        os.unlink(pidfile)
    except subprocess.CalledProcessError as e:
        sys.stderr.write("Unable to kill: %s.\n" % e)

    if os.path.exists(dest_data) or os.path.exists(dest_log):
        sys.stderr.write("Destination data or log file already exists. Not copying.\n")
        sys.exit(1)

    if not os.path.exists(dest_dir):
        try:
            os.mkdir(dest_dir)
        except OSError as e:
            sys.stderr.write("Unable to create destination directory: %s.\n" % e)
            sys.exit(1)

    time.sleep(0.5)

    shutil.copyfile(logfile, dest_log)
    try:
        if os.path.exists(datafile):
            subprocess.check_call("gzip -c %s > %s" % (datafile, dest_data), shell=True)
    except subprocess.CalledProcessError as e:
        sys.stderr.write("Unable to compress data file: %s.\n" % e)

    try:
        subprocess.check_call("sudo rm -f %s %s" % (datafile, logfile), shell=True)
    except subprocess.CalledProcessError as e:
        sys.stderr.write("Unable to remove data and log files: %s.\n" % e)




if __name__ == "__main__":
    options,args = parser.parse_args()
    if len(args) != 2:
        parser.error("Need action and filename.")
    action,filename = args

    filename = "%s.%s" % (filename, socket.gethostname())

    if action == 'start':
        start_tcpdump(filename, iface=options.INTERFACE, size=options.SIZE)
    elif action == 'stop':
        stop_tcpdump(filename, dest_dir=options.DEST_DIR)
    else:
        parser.error("Unknown action: %s." % action)
