#!/usr/bin/env python
# -*- coding: utf-8 -*-
# ======================================================
# Docker configuration and control on OpenFW UTM Community
# ======================================================

import os
import sys
import subprocess
import re

SETTINGS_FILE = "/var/efw/docker/settings"
START_LOCAL = "/var/efw/inithooks/start.local"


def log(message):
    print "[docker-init] {}".format(message)


def run_command(command):
    """Executes a system command (equivalent to Bash behavior)."""
    # Using subprocess.Popen for Python 2.7 compatibility to capture output/errors safely
    process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    process.communicate()
    return process.returncode


# ------------------------------------------------------
# Load settings
# ------------------------------------------------------
if not os.path.isfile(SETTINGS_FILE):
    log("Configuration file {} not found!".format(SETTINGS_FILE))
    sys.exit(1)

# Simulates Bash 'source' by reading key=value pairs
settings = {}
with open(SETTINGS_FILE, "r") as f:
    for line in f:
        line = line.strip()
        if line and not line.startswith("#"):
            if "=" in line:
                key, val = line.split("=", 1)
                # Remove quotes if present
                settings[key.strip()] = val.strip().strip('"').strip("'")

# Validate mandatory variables
if "DOCKER_ENABLED" not in settings or not settings["DOCKER_ENABLED"]:
    log("Error: DOCKER_ENABLED not defined in settings!")
    sys.exit(1)

DOCKER_ENABLED = settings["DOCKER_ENABLED"]


def apply_docker_firewall():
    config_path = "/var/efw/dockerfw/config"
    chain = "DOCKER-USER"

    # Ensure chain exists
    run_command("iptables -N {}".format(chain))

    # Flush DOCKER-USER chain
    run_command("iptables -F {}".format(chain))

    # Flush DOCKER chain to prevent rules outside DOCKER-USER chain
    run_command("iptables -F DOCKER")

    # Mandatory rule at the top
    run_command("iptables -I {} 1 -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT".format(chain))

    # If destination is not docker networks → ignore
    run_command("iptables -I {} 2 ! -d 100.125.0.0/16 -j RETURN".format(chain))
    run_command("iptables -I {} 3 ! -d 100.126.0.0/16 -j RETURN".format(chain))

    line_num = 2

    if not os.path.isfile(config_path):
        # If firewall config file does not exist, apply final drops only
        run_command("iptables -A {} -j NFLOG --nflog-prefix 'DOCKER:REJECT'".format(chain))
        run_command("iptables -A {} -j REJECT".format(chain))
        return

    with open(config_path, "r") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            
            # Split by commas (IFS=',')
            parts = line.split(",")
            if len(parts) < 7:
                continue

            status, proto, src, ports, action, obs, logflag = [p.strip() for p in parts[:7]]

            if status != "on":
                continue

            # Protocols
            if proto == "any":
                protos = ["tcp", "udp"]
            else:
                protos = proto.replace("&", " ").split()

            # Sources (If empty/blank, fallback to 0.0.0.0/0)
            if not src:
                srcs = ["0.0.0.0/0"]
            else:
                srcs = src.replace("&", " ").split()

            # Ports
            ports_list = ports.replace("&", " ").split()

            for p in protos:
                for s in srcs:
                    for d in ports_list:
                        rule = "-p {} --dport {}".format(p, d)

                        # Do not append -s if source is explicitly 0.0.0.0 or 0.0.0.0/0 (matches any)
                        if s != "0.0.0.0" and s != "0.0.0.0/0":
                            rule = "{} -s {}".format(rule, s)

                        if logflag == "on":
                            run_command("iptables -I {} {} {} -j NFLOG --nflog-prefix 'DOCKER:{}'".format(chain, line_num, rule, action))
                            line_num += 1

                        run_command("iptables -I {} {} {} -j {}".format(chain, line_num, rule, action))
                        line_num += 1

    # Final general drops NFLOG
    run_command("iptables -A {} -j NFLOG --nflog-prefix 'DOCKER:REJECT'".format(chain))

    # Mandatory final REJECT
    run_command("iptables -A {} -j REJECT".format(chain))


# ------------------------------------------------------
# Function: Controls boot startup
# ------------------------------------------------------
def configure_startup(state):
    log("Configuring Docker startup on boot...")

    # If file does not exist, create empty
    if not os.path.isfile(START_LOCAL):
        log("{} does not exist, creating...".format(START_LOCAL))
        with open(START_LOCAL, "w") as f:
            f.write("#!/bin/bash\n")
        os.chmod(START_LOCAL, 0o755)

    # Ensure shebang is correct
    with open(START_LOCAL, "r") as f:
        lines = f.readlines()

    if not lines or not re.match(r"^#!/bin/bash", lines[0]):
        lines.insert(0, "#!/bin/bash\n")
        with open(START_LOCAL, "w") as f:
            f.writelines(lines)
        os.chmod(START_LOCAL, 0o755)

    # Remove old lines containing 'docker' (equivalent to sed -i '/docker/d')
    lines = [line for line in lines if "docker" not in line]

    if state == "on":
        lines.append("/etc/init.d/docker start\n")

    with open(START_LOCAL, "w") as f:
        f.writelines(lines)


def apply_docker_state():
    if DOCKER_ENABLED == "on":
        log("Docker enabled — starting service...")
        # cat /etc/docker/daemon.json.tmpl > /etc/docker/daemon.json
        if os.path.isfile("/etc/docker/daemon.json.tmpl"):
            run_command("cat /etc/docker/daemon.json.tmpl > /etc/docker/daemon.json")
        
        run_command("/etc/init.d/docker restart")
        configure_startup("on")
        run_command("monit monitor dockerd")
        run_command("/etc/init.d/monit reload")
        
    elif DOCKER_ENABLED == "off":
        log("Docker disabled — stopping service...")
        run_command("/etc/init.d/docker stop")
        configure_startup("off")
        run_command("monit unmonitor dockerd")
        run_command("/etc/init.d/monit reload")
        
    else:
        log("Invalid value in DOCKER_ENABLED: {} (use on/off)".format(DOCKER_ENABLED))
        sys.exit(1)


# ------------------------------------------------------
# Main Execution
# ------------------------------------------------------
if len(sys.argv) > 1 and sys.argv[1] == "--firewall":
    apply_docker_firewall()
    sys.exit(0)

apply_docker_state()
apply_docker_firewall()