#!/usr/bin/python3
#
# docker-proxy-wrapper
#
# This wrapper script handles addition and removal of a CIS compliant iptables INPUT rule for a given docker-proxy
# invocation, it will try to add the rule, then run docker-proxy with the original arguments, and when receiving a SIGINT
# it will terminate docker-proxy and then remove the iptables rule (if it was added)
#
# For some reason I can't quite figure out, docker-proxy feels the need to communicate success/failure as (0|1)\n over
# a pipe on fd 3, so we need to relay that information from docker-proxy to the caller too, but since process.Popen
# will inherit the fd, we shouldn't need to do anything extra to achieve that
#
# As a general note, keep in mind that normal off-host traffic targeting containers does NOT pass through docker-proxy
# or the INPUT chain, but is rather DNATed to the docker0 bridge in the nat table's PREROUTING chain, docker-proxy
# is only used for loopback traffic on the host, specifically in our case that is almost always only xray traces
# In theory we may in fact be able to forgo docker-proxy in all cases except for xray, but it's very hard to confirm/prove
# that it would not break corner cases
#
#------------------------------------------------
# Usage of docker-proxy:
#   -container-ip string
#         container ip
#   -container-port int
#         container port (default -1)
#   -host-ip string
#         host ip
#   -host-port int
#         host port (default -1)
#   -proto string
#         proxy protocol (default "tcp")
#------------------------------------------------

import os
import re
import sys
import locale
import signal
import subprocess
from functools import partial
from pprint import pprint

# Parse the arguments
args = dict(zip(sys.argv[1::2], sys.argv[2::2]))


ALLOWED_ORIGINS = [ '172.17.0.0/16', '127.0.0.0/8' ]
# Parse extra origins from the environment
try:
    if 'PROXY_ALLOWED_ORIGINS' in os.environ:
        ALLOWED_ORIGINS += re.split('\s+', os.environ['PROXY_ALLOWED_ORIGINS'])
        # Make sure our list is unique
        ALLOWED_ORIGINS = list(set(ALLOWED_ORIGINS))
except Exception as e:
    print('Encountered Exception while parsing PROXY_ALLOWED_ORIGINS from environment:', file=sys.stderr)
    pprint(e, stream=sys.stderr)
    pass

# Step one: Add our new INPUT RULE
try:
    # Find our INPUT anchor from 'iptables -S INPUT'
    # NOTE: iptables -S outputs the policy as line 1, but that's ok, since the list is 0 based for python, but 1 based for iptables
    rules = subprocess.run(['iptables', '-S', 'INPUT'], text=True, check=True, capture_output=True).stdout
    rules = rules.splitlines(False)

    # Find our anchor rule's index
    pos = rules.index('-A INPUT -m comment --comment "DOCKER-PROXY-WRAPPER::ANCHOR::DO-NOT-DELETE"')

    # Insert our new INPUT rule immediately after the anchor rule
    # NOTE: There is a race condition here if someone else modifies the INPUT table above our anchor our position will be wrong.
    #       The *correct* solution would be to create and use a dedicated name user chain, but since AWS Inspector is brain dead
    #       that would lead to reported non-compliance, so we just have to take the (small) risk
    # *IMPORTANT* The order of the parameters for this rule *do* matter because AWS Inspector is DUMB, see exhibit (a) above
    for src in ALLOWED_ORIGINS:
        cmd =['iptables', '-I', 'INPUT', str(pos+1), '-s', src, '-p', args['-proto'], '-m', args['-proto'], '--dport', args['-host-port'], '-m', 'state', '--state', 'NEW', '-j', 'ACCEPT']
        print('Adding docker-proxy iptables rule: ' + ' '.join(cmd))
        result = subprocess.run(cmd, text=True, check=True, capture_output=True)
        rule_was_added = True

except Exception as e:
    rule_was_added = False
    print('Encountered Exception while installing iptables rule:', file=sys.stderr)
    pprint(e, stream=sys.stderr)
    pass


# Step two: Run the actual docker-proxy command

# Simple wrapper to prop
def signaled(proxy, signo, frame):
    proxy.send_signal(signo)

try:
    cmd = sys.argv
    cmd[0] = 'docker-proxy'
    proxy = subprocess.Popen(cmd, close_fds=False, shell=False)
    for sig in [signal.SIGINT, signal.SIGTERM]:
        signal.signal(sig,  lambda signo, frame: signaled(proxy, signo, frame))
    print(f'Running docker-proxy worker as pid {proxy.pid}: ' + ' '.join(cmd))
    proxy.wait()
except Exception as e:
    print('Encountered Exception while running docker-proxy:', file=sys.stderr)
    pprint(e, stream=sys.stderr)
    pass

# Step three: Remove the INPUT rule (if it was added)
if rule_was_added:
    try:
        # See the rant about in step one about why this is far from the sanest or safest way to manage the rules, but... AWS Inspector stupidity
        # Fortunately we can at least rely on iptables to just match the rule by it's parameters though, so it's not nearly as risky as deleting by offset
        for src in ALLOWED_ORIGINS:
            cmd =['iptables', '-D', 'INPUT', '-s', src, '-p', args['-proto'], '-m', args['-proto'], '--dport', args['-host-port'], '-m', 'state', '--state', 'NEW', '-j', 'ACCEPT']
            print('Removing docker-proxy iptables rule: ' + ' '.join(cmd))
            result = subprocess.run(cmd, check=True, capture_output=True)
            rule_was_removed = True
    except Exception as  e:
        rule_was_removed = False
        print('Encountered Exception while removing iptables rule:', file=sys.stderr)
        pprint(e, stream=sys.stderr)
        pass





