#!/usr/bin/env python3
from docker import Client
from argparse import ArgumentParser
from sys import stderr
from os import getuid

from dnc.pipeline import Pipeline, LoopbackProducer, InterfaceProducer, JsonConsumer, StdoutConsumer
from dnc import CaptureFlags, ContainerError, PrivilegeError


def _main(container, bridge="docker0", capture=CaptureFlags.Egress, docker='', **kwargs):
    """ Initialize and start packet capture for `container` on specific `bridge`

    :param container: str, container name or id
    :param bridge: str, name of the docker bridge
    :param capture: CaptureFlags, what kind of traffic to capture
    :param docker: str, url to the docker socket
    :param kwargs: dict, additional arguments passed to the pipeline
    :return:
    """
    root = 'root' in kwargs and kwargs['root']
    if root and getuid() != 0:
        raise PrivilegeError("Execution requires root (uid: {})".format(getuid()))

    cli = Client(base_url=docker)
    cont = cli.inspect_container(container)
    ns = cont.get("NetworkSettings", None)
    if not ns:
        raise ContainerError("Container has no associated network settings")
    mac = ns.get("MacAddress", None)
    if not mac:
        raise ContainerError("Container has no MAC Address")
    state = cont.get("State", None)
    if not state:
        raise ContainerError("Container has no state")
    pid = state.get("Pid", None)
    if not pid:
        raise ContainerError("Container has no PID (is it running?)")
    ip = ns.get("IPAddress", None)
    if not pid:
        raise ContainerError("Container has no IP address")

    # check invocation flags
    verbose = 'verbose' in kwargs and kwargs['verbose']
    jsonout = 'json' in kwargs and kwargs['json']

    # print out container bridge, name and ID when we're verbose
    if verbose:
        name, cid = cont.get("Name", "/unnamed"), cont.get("Id", "(no-id)")
        msg = '[+] Initializing packet capture on: "{b}" for container: "{c}@{h}"'
        print(msg.format(b=args.bridge, c=name[1:], h=cid), file=stderr)

    # depending on the output, instantiate a consumer
    if jsonout:
        consumer = JsonConsumer(ip)
    else:
        consumer = StdoutConsumer(ip)

    # create pipeline, add producers and consumers
    pipeline = Pipeline(**kwargs)
    pipeline.add_producer(LoopbackProducer(pid, capture))
    pipeline.add_producer(InterfaceProducer(bridge, mac, capture))
    pipeline.add_consumer(consumer)

    try:
        # execute the pipeline and loop until we're done
        pipeline.run()
    except KeyboardInterrupt:
        if verbose:
            print("\b\bExiting", file=stderr)
        pipeline.running = False
        if jsonout:
            import json
            print(json.dumps(consumer.buffered, indent=3))


def parse_args():
    """ Parse command line arguments and return them """
    arguments = ArgumentParser()
    arguments.add_argument('--verbose', default=False, action='store_true')
    arguments.add_argument('--no-root', default=False, action='store_true')
    arguments.add_argument('-b', '--bridge', default='docker0', help='Docker bridge to use')
    arguments.add_argument('-c', '--capture', default='ingress,egress', help='Direction of traffic to capture [egress,ingress]')
    arguments.add_argument('-j', '--json', default=False, help='buffer data and output JSON at the end', action='store_true')
    arguments.add_argument('-d', '--docker', default='unix://var/run/docker.sock', help='')
    arguments.add_argument('container')

    # parse arguments
    arg = arguments.parse_args()
    capture_args = arg.capture.lower()

    # parse traffic direction spec
    cap = CaptureFlags.Invalid
    if 'ingress' in capture_args:
        cap |= CaptureFlags.Ingress
    if 'egress' in capture_args:
        cap |= CaptureFlags.Egress
    arg.cap = cap

    return arg


if __name__ == '__main__':
    args = parse_args()
    _main(args.container, bridge=args.bridge, docker=args.docker,
          capture=args.cap, verbose=args.verbose, json=args.json,
          root=not args.no_root)
