import os
import re
from subprocess import check_output, call


def get_ips():
    lines = check_output(['ip', 'addr', 'show']).split("\n")
    ips = []
    for line in lines:
        aa=re.match(r"inet (\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})",line.strip())
        if aa:
            ip = aa.group(1)
            if ip != "127.0.0.1":
                ips.append(ip)
                user = ip.replace(".","")
    return ips

def get_users():
    lines = check_output(['cat','/etc/passwd']).split("\n")
    users = []
    for line in lines:
        aa=re.match(r"(ip\d+).+", line.strip())
        if aa:
            users.append(aa.group(1))
    return users

def create_users(count):
    for i in range(1,count+1):
        os.system("useradd -M -N ip{}".format(i))
        os.system("echo \"prx\nprx\n\" | passwd ip{}".format(i))


def print_ips(ips):
    lines = []
    for i in range(0,len(ips)):
        lines.append("IP{}=\"{}\"".format(i, ips[i]))
    return "\n".join(lines)

def print_marks(ips, users):
    lines = []
    for i in range(0, len(ips)):
        lines.append("$IPT -t mangle -A OUTPUT -m owner --uid-owner {} -j MARK --set-mark 0x{}".format(users[i], i+1))
    return "\n".join(lines)

def print_trans(ips):
    lines = []
    for i in range(0, len(ips)):
        lines.append("$IPT -t nat -A POSTROUTING -m mark --mark 0x{} -j SNAT --to-source $IP{}".format(i+1, i))
    return "\n".join(lines)

def main():
    ips = get_ips()
    users = get_users()
    if len(ips) > len(users):
        create_users(len(ips))
        users = get_users()


    pips = print_ips(ips)
    marks = print_marks(ips, users)
    trans = print_trans(ips)
    print_conf(pips, marks, trans)


def print_conf(pips, marks, trans):
    s = """
#!/bin/bash

IPT="/sbin/iptables"

{}

# Flush old rules
$IPT -F
$IPT -F -t nat
$IPT -F -t mangle
$IPT -X
$IPT -X -t nat
$IPT -X -t mangle

# Set default policies
$IPT -P INPUT ACCEPT
$IPT -P FORWARD ACCEPT
$IPT -P OUTPUT ACCEPT

# Mark packages
{}

# Accept all traffic through the loopback interface
$IPT -A INPUT -i lo -j ACCEPT

# Accept traffic of established connetctions
$IPT -A INPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT

# Accept safe ICMP requests
$IPT -A INPUT -p icmp --icmp-type 3 -j ACCEPT
$IPT -A INPUT -p icmp --icmp-type 8 -j ACCEPT
$IPT -A INPUT -p icmp --icmp-type 12 -j ACCEPT

# Accept SSH
$IPT -A INPUT -p tcp -m conntrack --ctstate NEW --dport 22 -j ACCEPT
$IPT -A INPUT -p tcp -m conntrack --ctstate NEW --dport 22101 -j ACCEPT
$IPT -A INPUT -p tcp -m conntrack --ctstate NEW --dport 22102 -j ACCEPT
$IPT -A INPUT -p tcp -m conntrack --ctstate NEW --dport 22103 -j ACCEPT
$IPT -A INPUT -p tcp -m conntrack --ctstate NEW --dport 22104 -j ACCEPT

# Public IP translation
{}
    """.format(pips,marks, trans)

    print(s)

if __name__ == '__main__':
    main()