#!/usr/bin/python3
# -*- coding: utf-8 -*-

import argparse
import os
import shutil
import sys
import tarfile
import tempfile

APK_WORLD = 'etc/apk/world'
APK_INSTALLED = 'lib/apk/db/installed'
APK_SCRIPTS = 'lib/apk/db/scripts.tar'
APK_TRIGGERS = 'lib/apk/db/triggers'

ETC_PASSWD = 'etc/passwd'
ETC_GROUP = 'etc/groups'
ETC_SHADOW = 'etc/shadow'

def makedirs(path, mode=0o755, uid=100000, gid=100000):
    try:
        os.mkdir(path, mode)
        os.chown(path, uid, gid)
    except FileNotFoundError:
        makedirs(os.path.dirname(path), mode, uid, gid)
        os.mkdir(path, mode)
        os.chown(path, uid, gid)
    except FileExistsError:
        pass

def merge_apk_world(layers):
    world = []
    for layer in layers:
        try:
            with open(os.path.join(layer, APK_WORLD), 'r') as f:
                for line in f:
                    if line not in world:
                        world.append(line)
        except:
            continue
    makedirs(os.path.join(layers[-1], os.path.dirname(APK_WORLD)))
    with open(os.path.join(layers[-1], APK_WORLD), 'w') as f:
        f.writelines(world)
    os.chown(os.path.join(layers[-1], APK_WORLD), 100000, 100000)

def merge_apk_installed(layers):
    installed = []
    for layer in layers:
        try:
            with open(os.path.join(layer, APK_INSTALLED), 'r') as f:
                buffer = []
                for line in f:
                    if line.startswith('C:'):
                        buffer = ''.join(buffer)
                        if buffer not in installed:
                            installed.append(buffer)
                        buffer = []
                    buffer.append(line)
                buffer = ''.join(buffer)
                if buffer not in installed:
                    installed.append(buffer)
        except:
            continue
    makedirs(os.path.join(layers[-1], os.path.dirname(APK_INSTALLED)))
    with open(os.path.join(layers[-1], APK_INSTALLED), 'w') as f:
        f.writelines(installed)
    os.chown(os.path.join(layers[-1], APK_INSTALLED), 100000, 100000)

def merge_apk_scripts(layers):
    tmp_tar_path = tempfile.mkstemp()[1]
    files_in_tar = []
    with tarfile.open(tmp_tar_path, 'w:') as tmp_tar:
        for layer in layers:
            tar_path = os.path.join(layer, APK_SCRIPTS)
            if os.path.exists(tar_path):
                with tarfile.open(tar_path, 'r:') as tar:
                    for member in tar.getmembers():
                        if member.name not in files_in_tar:
                            buffer = tar.extractfile(member)
                            tmp_tar.addfile(member, buffer)
                            files_in_tar.append(member.name)
    if files_in_tar:
        makedirs(os.path.join(layers[-1], os.path.dirname(APK_SCRIPTS)))
        shutil.move(tmp_tar_path, os.path.join(layers[-1], APK_SCRIPTS))
        os.chown(os.path.join(layers[-1], APK_SCRIPTS), 100000, 100000)
    else:
        os.unlink(tmp_tar_path)

def merge_apk_triggers(layers):
    triggers = []
    for layer in layers:
        try:
            with open(os.path.join(layer, APK_TRIGGERS), 'r') as f:
                for line in f:
                    if line not in triggers:
                        triggers.append(line)
        except:
            continue
    makedirs(os.path.join(layers[-1], os.path.dirname(APK_TRIGGERS)))
    with open(os.path.join(layers[-1], APK_TRIGGERS), 'w') as f:
        f.writelines(triggers)
    os.chown(os.path.join(layers[-1], APK_TRIGGERS), 100000, 100000)

def merge_etc_passwd(layers):
    passwd = {}
    for layer in layers:
        try:
            with open(os.path.join(layer, ETC_PASSWD), 'r') as f:
                for line in f:
                    passwd[line.split(':')[0]] = line
        except:
            continue
    makedirs(os.path.join(layers[-1], os.path.dirname(ETC_PASSWD)))
    with open(os.path.join(layers[-1], ETC_PASSWD), 'w') as f:
        f.writelines(passwd.values())
    os.chown(os.path.join(layers[-1], ETC_PASSWD), 100000, 100000)

def merge_etc_group(layers):
    groups = {}
    for layer in layers:
        try:
            with open(os.path.join(layer, ETC_GROUP), 'r') as f:
                for line in f:
                    name,pwd,gid,users = line.split(':')
                    name = splitline[0]
                    users = splitline[3].strip().split(',')
                    if name not in groups:
                        groups[name] = [name,pwd,gid,users]
                    else:
                        groups[name][1] = pwd
                        groups[name][2] = gid
                        for user in users:
                            if user not in groups[name][3]:
                                groups[name][3].append(user)
        except:
            continue
    for group in groups.values():
        group[3] = '{}\n'.format(','.join(group[3]))
    makedirs(os.path.join(layers[-1], os.path.dirname(ETC_GROUP)))
    with open(os.path.join(layers[-1], ETC_GROUP), 'w') as f:
        f.writelines([':'.join(group) for group in groups.values()])
    os.chown(os.path.join(layers[-1], ETC_GROUP), 100000, 100000)

def merge_etc_shadow(layers):
    shadow = {}
    for layer in layers:
        try:
            with open(os.path.join(layer, ETC_SHADOW), 'r') as f:
                for line in f:
                    shadow[line.split(':')[0]] = line
        except:
            continue
    makedirs(os.path.join(layers[-1], os.path.dirname(ETC_SHADOW)))
    with open(os.path.join(layers[-1], ETC_SHADOW), 'w') as f:
        f.writelines(shadow.values())
    os.chown(os.path.join(layers[-1], ETC_SHADOW), 100000, 100042)


parser = argparse.ArgumentParser(description='APK database merge script')
parser.add_argument('layers', help='Path to LXC layers to be merged', nargs=argparse.REMAINDER)

if len(sys.argv) < 3:
    parser.print_usage()
    sys.exit(1)
args = parser.parse_args()

merge_apk_world(args.layers)
merge_apk_installed(args.layers)
merge_apk_scripts(args.layers)
merge_apk_triggers(args.layers)

merge_etc_passwd(args.layers)
merge_etc_group(args.layers)
merge_etc_shadow(args.layers)