import urllib3
urllib3.disable_warnings()
import pymysql
import subprocess
import os
import random

REGISTRY_URL = 'localhost:5000'


def make_empty_main_py():
    containers = get_container_names()
    for c in containers:
        container_name = c['name']
        f = open("{}/_main.py".format(container_name), 'w')
        f.write("print 'empty!'")
        f.close()

def get_team_names():
    conn = pymysql.connect(host='localhost', user='root', password='zp54cprfd7wokhwg',
        db='ctf', charset='utf8mb4',
        cursorclass=pymysql.cursors.DictCursor)
    curs = conn.cursor()
    query = """select id, team_name, password from teams"""
    curs.execute(query)
    rows = curs.fetchall()
    conn.close()
    return rows

def get_container_names():
    conn = pymysql.connect(host='localhost', user='root', password='zp54cprfd7wokhwg',
        db='ctf', charset='utf8mb4',
        cursorclass=pymysql.cursors.DictCursor)
    curs = conn.cursor()
    query = """select name from services"""
    curs.execute(query)
    rows = curs.fetchall()
    conn.close()
    return rows

def build_image(image, container_config_dir):
    command = ["sudo","docker", "build", "-t", image["path"],
                os.path.join(container_config_dir, image["name"])]
    process = subprocess.Popen(command, stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE)
    stdout, stderr = process.communicate() 
    if process.returncode != 0:
        print "Something went wrong when building image of %s" % (image['name'])
        print "Stdout: %s" % (stdout)
        print "Stderr: %s" % (stderr)
    else:
        print "Built image for service %s with tag %s" % (image['name'], image['name'])

def team_login(team):
    command = ["sudo", "docker", "login", "--username={}".format(team['team_name']), "--password={}".format(team['password']), "localhost:5000"]
    process = subprocess.Popen(command, stdout=subprocess.PIPE,
                stderr=subprocess.PIPE)
    stdout, stderr = process.communicate() 
    if process.returncode != 0:
        print "login fail"
        print "Stdout: %s" % (stdout)
        print "Stderr: %s" % (stderr)
        return False
    else:
        print "login success %s" %(team['team_name'])
        return True

def push_image(image, team):
    if team_login(team):
        command = ["sudo", "docker", "push", image['path']]
        process = subprocess.Popen(command, stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE)    
        stdout, stderr = process.communicate() 
        print "Stdout: %s" % (stdout)
        if process.returncode != 0:
            print "push image"
            print "Stderr: %s" % (stderr)
        else:
            print "push success %s" %(image['path'])

def toggle_main_py(container_name):
    command = ["mv", "./{}/main.py".format(container_name), "./{}/tmp.py".format(container_name), 
        "&&", "mv", "./{}/_main.py".format(container_name), "./{}/main.py".format(container_name),
        "&&", "mv", "./{}/tmp.py".format(container_name), "./{}/_main.py".format(container_name)]
    process = subprocess.Popen(command, stdout=subprocess.PIPE,
                stderr=subprocess.PIPE) 
    stdout, stderr = process.communicate()
    if process.returncode != 0:
        print "name change fail"
        print "Stdout: %s" % (stdout)
        print "Stderr: %s" % (stderr)
    else:
        print "main.py file swap success %s" %(container_name)

def init():
    make_empty_main_py()

def main():
    teams = get_team_names()
    containers = get_container_names()
    for t in teams:
        for c in containers:
            namespace = t['team_name']
            container_name = c['name']
            type_ = "exploit"
            if bool(random.getrandbits(1)):
                toggle_main_py(container_name)
            image = { 'path' : "{}/{}/{}_{}".format(REGISTRY_URL , namespace, type_,container_name), "name": container_name}
            build_image(image, '.')
            if bool(random.getrandbits(1)):
                push_image(image, t)
    

if __name__ == '__main__':
    main()