#!/usr/bin/python3

import sys
import pathlib
import os
import argparse

def set_idmap_in_config(config_location, base_map, dry_run=False):
    config_data = ""

    found_user = False
    found_group = False
    with open(config_location, "r") as File:
        for row in File:
            if row.startswith("lxc.idmap"):
                data = row.split("=")
                conf_mapping = data[1].split(" ")
                if conf_mapping[0] == "u":
                    config_data += "lxc.idmap = u 0 {} 65536\n".format(base_map)
                    found_user = True
                elif conf_mapping[1] == "g":
                    config_data += "lxc.idmap = g 0 {} 65536\n".format(base_map)
                    found_group = True
            else:
                config_data += row

    if not found_user:
        config_data += "lxc.idmap = u 0 {} 65536\n".format(base_map)
    if not found_group:
        config_data += "lxc.idmap = g 0 {} 65536\n".format(base_map)

    if not dry_run:
        with open(config_location, "w") as File:
            File.write(config_data)
    else:
        print("New Config")
        print(config_data)

def map_user_ids(name, base_map, uid_map, dry_run=False):
    start_path = pathlib.Path("/var/lib/lxc", name, "rootfs")
    path_list = list(start_path.rglob("*"))
    path_list.append(start_path)
    path_list.append(pathlib.Path("/var/lib/lxc", name))

    new_base_map = base_map + uid_map
    for i in path_list:
        stat = i.lstat()
        local_uid = stat.st_uid
        local_gid = stat.st_gid

        if local_uid >= new_base_map * 2:
            local_uid = local_uid % 65536
            local_gid = local_gid % 65536

        if not dry_run:
            os.lchown(i, local_uid + new_base_map, local_gid + new_base_map)
        else:
            print("chowning {} from {} to {}".format(i, local_uid, local_uid + new_base_map))

    set_idmap_in_config(pathlib.Path("/var/lib/lxc", name, "config"), new_base_map, dry_run)

def main():
    parser = argparse.ArgumentParser(
            prog="User ID mapping",
            description="maps base uids to a higher number")

    parser.add_argument("server_name")
    parser.add_argument("mapping")
    parser.add_argument("-b", "--base-map")
    parser.add_argument("-m", "--multiplier")
    parser.add_argument("-d", "--dry-run", action=argparse.BooleanOptionalAction)

    args = parser.parse_args()

    base_map = 900000000
    if args.base_map:
        base_map = int(args.base_map)

    multiplier = 100000
    if args.multiplier:
        multiplier = int(args.multiplier)

    server_name = args.server_name
    mapping = int(args.mapping)

    print(args.dry_run)

    map_user_ids(server_name, base_map, multiplier * mapping, args.dry_run)

if __name__ == "__main__":
    main()