import os
import sys
import time
import json


if len(sys.argv) < 2:
    print("usage: script type[dump]")
    exit(0)
    pass

if sys.argv[1] != "check-config":
    import MySQLdb
    pass

check_next_gen_priv = False
global_pyconnection = None
pyconnection_env = "dev"
pyconnection_db_pass = None
if len(sys.argv) > 4:
    print("sys.argv[2]: " + sys.argv[2])
    print("sys.argv[4]: " + sys.argv[4])
    if sys.argv[2] != "":
        pyconnection_env = sys.argv[2]
    if pyconnection_env == "dev" or pyconnection_env == "staging":
        check_next_gen_priv = True
    pyconnection_db_pass = sys.argv[3]
    if sys.argv[4] == "true":
        check_next_gen_priv = False
if check_next_gen_priv:
    import pymysql
    pass

cur_dir = os.path.split(os.path.realpath(__file__))[0]
config_file="/etc/dbaas/config.yaml"


env_param = os.getenv("ENV", default="dev")

if env_param == "local-staging":
    config_file = "/etc/dbaas/config-staging.yaml"

print("config_file is " + config_file)

dump_rolebinding_file = cur_dir + "/ALL_ROLEBINDING.json"
dump_privilege_file = cur_dir + "/ALL_PRIVILEGE.json"

global_connection = None


def read_account_config():
    import yaml
    f = open(config_file, 'r', encoding='utf-8')
    cfg = f.read()
    d = yaml.safe_load(cfg)
    return d

def get_mysql_connection():
    import MySQLdb
    global global_connection
    config_info = read_account_config()
    mysql_conf = config_info["mysql"]
    if global_connection != None:
        return global_connection
    global_connection = MySQLdb.connect(
        host=mysql_conf["host"],
        port=mysql_conf["port"],
        user=mysql_conf["user"],
        password=mysql_conf["password"],
        database=mysql_conf["database"]
        )
    return global_connection

def run_sql(conn, sql):
    with conn.cursor() as cursor:
        cursor.execute(sql)
        conn.commit()
        pass
    return
def get_one_sql(conn, sql):
    results = get_sql(conn, sql)
    if len(results) != 1:
        raise Exception("result size is not 1, " + str(len(results)))
    return results[0]
def get_sql(conn, sql):
    results = []
    with conn.cursor() as cursor:
        cursor.execute(sql)
        results = cursor.fetchall()
        conn.commit()
    return results

def read_json_from_file(filename):
    f = open(filename)
    data = json.load(f)
    f.close()
    return data

def write_content_to_file(filename, content):
    file = open(filename,"w")
    file.write(content)
    file.close()
    return
    
def get_system_list(dump_rolebinding_file):
    rolebinding_info = read_json_from_file(dump_rolebinding_file)
    return rolebinding_info.keys()

def dump_config():
    rolebinding_configs = get_online_rbac_rolebinding() 
    write_content_to_file(dump_rolebinding_file, json.dumps(rolebinding_configs, indent=1, sort_keys=True))
    privilege_configs = get_online_rbac_privilege()
    write_content_to_file(dump_privilege_file, json.dumps(privilege_configs, indent=1, sort_keys=True))
    return

def split_dump_data():
    output_dir = cur_dir + "/output"
    os.makedirs(name=output_dir, exist_ok=True)
    for system, rbac_info in rbac_configs.items():
        output_file = output_dir + "/" + system + ".role_binding"
        write_content_to_file(output_file, json.dumps(rbac_info, indent=1))
        pass
    return
def get_online_rbac_privilege():
    ret_dict = {}
    connection = get_mysql_connection()
    sql = "select resource_uri,url,`system`,description from rbac_privileges order by `system`,resource_uri"
    rows = get_sql(connection, sql)
    for row in rows:
        resource_uri = row[0]
        url = row[1]
        system = row[2]
        description = row[3]
        if system not in ret_dict:
            ret_dict[system] = {}
            pass
        ret_dict[system][resource_uri] = {"url": url, "description":description}
    return ret_dict

def get_online_rbac_rolebinding():
    ret_dict = {}
    connection = get_mysql_connection()
    sql = "select rbac_roles.name, rbac_privileges.resource_uri, rbac_privileges.system from rbac_role_privileges, rbac_roles, rbac_privileges where rbac_roles.id=rbac_role_privileges.rbac_role_id and rbac_privileges.id=rbac_role_privileges.rbac_privilege_id order by rbac_privileges.resource_uri"
    rows = get_sql(connection, sql)
    for row in rows:
        role_name = row[0]
        uri = row[1]
        system = row[2]
        if system not in ret_dict:
            ret_dict[system] = {}
            pass
        if role_name not in ret_dict[system]:
            ret_dict[system][role_name] = []
            pass
        ret_dict[system][role_name].append(uri)
    return ret_dict
        
def get_not_in_target_list(orig_list, target_list):
    ret_list = []
    for item in orig_list:
        if not item in target_list:
            ret_list.append(item)
            pass
        pass
    return ret_list

def get_privilige_items_by_add_list(add_list, privilege_dict):
    ret_dict = {}
    not_in_list = []
    if len(add_list) != 0:
        for item in add_list:
            uri_list = item["uri_list"]
            for uri in uri_list:
                if uri not in privilege_dict:
                    not_in_list.append(uri)
                else:
                    ret_dict[uri] = privilege_dict[uri]
                pass
            pass
        pass
    return ret_dict, not_in_list

def update_privilege(system, uri, url, description, is_preview):
    try:
        privilege_info = get_privilege_record(uri)
    except Exception:
        privilege_info = None
        pass
    if privilege_info != None:
        print("privilege_info has been already existed, skip " + uri)
        return
    
    sql = "insert into rbac_privileges(resource_uri,url,`system`,description) values('"+uri+"', '"+url+"', '"+system+"', '" + description + "')"
    print(sql)
    if is_preview:
        return
    connection = get_mysql_connection()
    run_sql(connection, sql)
    return

def get_role_record(role):
    sql = "select id, name from rbac_roles where name='"+role+"'"
    conn = get_mysql_connection()
    row = get_one_sql(conn, sql)
    return {"id": row[0], "name": row[1]}
    
def get_privilege_record(uri):
    sql = "select id, resource_uri, `system` from rbac_privileges where resource_uri=" +repr(uri)
    conn = get_mysql_connection()
    row = get_one_sql(conn, sql)
    return {"id": row[0], "resource_uri": row[1], "system": row[2]}

def add_rolebinding_record(role, uri, is_preview):
    role_info = get_role_record(role)
    privilege_info = get_privilege_record(uri)
    sql = "insert into rbac_role_privileges(rbac_role_id,rbac_privilege_id,created_at,updated_at) values("+str(role_info["id"])+","+str(privilege_info["id"])+",now(),now())"
    print(sql)
    if is_preview:
        return
    connection = get_mysql_connection()
    return run_sql(connection, sql)

def delete_rolebinding_record(role, uri, is_preview):
    role_info = get_role_record(role)
    try:
        privilege_info = get_privilege_record(uri)
    except Exception:
        print("no uri " + uri + " found , skip")
        return
    sql = "delete from rbac_role_privileges where rbac_role_id=" + str(role_info["id"]) + " and rbac_privilege_id=" + str(privilege_info["id"])
    print(sql)
    if is_preview:
        return
    connection = get_mysql_connection()
    return run_sql(connection, sql)
def update_config(system, dump_rolebinding_file, dump_privilege_file, is_preview):
    delete_list, add_list = compare_config(system, dump_rolebinding_file)

    all_privilege_dict = read_json_from_file(dump_privilege_file)
    privilege_dict = all_privilege_dict.get(system, {})
    privelige_items, not_in_list = get_privilige_items_by_add_list(add_list, privilege_dict)
    if len(not_in_list) != 0:
        print("update_config failed. some privilege infos are not included in ALL_PRIVILEGE file.")
        print(not_in_list)
        raise Exception("check privilege failed.")
    
    for item in add_list:
        uri_list = item["uri_list"]
        for uri in uri_list:
            try:
                privilege_info = get_privilege_record(uri)
            except Exception:
                privilege_info = None
                pass
            if privilege_info == None:
                continue
            if privilege_info['system'] != system:
                print("update config failed. privilege uri " + uri + " is configured in another system " + privilege_info["system"])
                raise Exception("check privilege_info failed.")
        pass

    # add privilege items
    for uri, privilege_info in privelige_items.items():
        update_privilege(system, uri, privilege_info["url"], privilege_info["description"], is_preview)
    
    # delete unused rolebindings
    for item in delete_list:
        role = item["role"]
        uri_list = item["uri_list"]
        for uri in uri_list:
            delete_rolebinding_record(role, uri, is_preview)
            pass
        pass

    # add new rolebindings
    for item in add_list:
        role = item["role"]
        uri_list = item["uri_list"]
        for uri in uri_list:
            add_rolebinding_record(role, uri, is_preview)
            pass
        pass

    return

def compare_rbac_rolebinding_with_online(system, compared_rolebinding_file):
    online_rbac_configs = get_online_rbac_rolebinding() 
    compared_rbac_configs = read_json_from_file(compared_rolebinding_file)
    online_configs = online_rbac_configs.get(system, {})
    compared_configs = compared_rbac_configs.get(system, {})
    return diff_system_configs(online_configs, compared_configs)

def compare_config(system, compared_rolebinding_file):
    delete_list, add_list = compare_rbac_rolebinding_with_online(system, compared_rolebinding_file)
    if len(delete_list) == 0  and len(add_list) == 0:
        print("NO DIFF")
        return delete_list, add_list
    print_compare_result(delete_list, add_list)
    return delete_list, add_list

def print_compare_result(delete_list, add_list):
    print("DELETE ROLE BINDINGS:")
    for item in delete_list:
        print(item["role"])
        print(item["uri_list"])
        print("==============================")
    print("\n\n")
    print("ADD ROLE BINDINGS:")
    for item in add_list:
        print(item["role"])
        print(item["uri_list"])
        print("==============================")
    print("\n\n")
    return

def diff_system_configs(online_configs, compared_configs):
    delete_list = []
    add_list = []
    for role, online_uri_list in online_configs.items():
        delete_items = get_not_in_target_list(online_uri_list, compared_configs.get(role,[]))
        if len(delete_items) != 0:
            delete_list.append({"role": role, "uri_list": delete_items})
            pass
        pass

    for role, compared_uri_list in compared_configs.items():
        add_items = get_not_in_target_list(compared_uri_list, online_configs.get(role,[]))
        if len(add_items) != 0:
            add_list.append({"role": role, "uri_list": add_items})
            pass
        pass
    return delete_list, add_list

def check_config__privilege_valid(system, rolebinding_items, privilege_info):
    privilege_items = privilege_info.get(system, {})
    for role, uri_list in rolebinding_items.items():
        for uri in uri_list:
            if uri not in privilege_items:
                err_msg = "[" + system + "] uri " + uri + " are not found in privilege_list"
                print(err_msg)
                raise Exception("check config failed. " + err_msg)
            pass
        pass
    return
def check_config__basic_org_role_config(system, rolebinding_items, privilege_info):
    org_level_roles = ["org:audit_admin", "org:billing_admin", "org:owner"]
    org_basic_roles = ["org:member"]
    check_config__check_role_issuperset(system, rolebinding_items, privilege_info, org_level_roles, org_basic_roles)
    org_level_roles = ["org:owner"]
    org_basic_roles = ["org:member", "org:billing_admin", "org:audit_admin", "project:owner", "project:dev"]
    check_config__check_role_issuperset(system, rolebinding_items, privilege_info, org_level_roles, org_basic_roles)
        

def check_config__org_member_config(system, rolebinding_items, privilege_info):
    url_privileges = rolebinding_items.get("org:member", [])
    for url_item in url_privileges:
        if "{project_id}" in url_item:
            err_msg = "[" + system + "] org:member includees invalid path: " + url_item 
            print(err_msg)
            raise Exception("check config failed. " + err_msg)
    return

def check_config__check_role_issuperset(system, rolebinding_items, privilege_info, super_roles, subset_roles):
    for org_basic_role in subset_roles:
        basic_privileges = rolebinding_items.get(org_basic_role, [])
        basic_set = set(basic_privileges)
        for org_role_item in super_roles:
            org_role_privileges = rolebinding_items.get(org_role_item, [])
            org_role_set = set(org_role_privileges)
            if not org_role_set.issuperset(basic_set):
                err_msg = "[" + system + "] org role " + org_role_item + " is not super set of " + org_basic_role
                print(err_msg)
                difference_set = basic_set.difference(org_role_set)
                print(difference_set)
                raise Exception("check config failed. " + err_msg)
            pass
        pass
    return

def check_config():
    print("============ check rbac config ============")

    # check priv - next gen priv must contains tidbs
    if check_next_gen_priv:
        print("start check next gen priv ...")
        check_config__next_gen_privs_valid()
        print("next gen priv check passed.")
    else:
        print("not enable check_next_gen_priv switch, ignore to check")

    rolebinding_info = read_json_from_file(dump_rolebinding_file)
    privilege_info = read_json_from_file(dump_privilege_file)
    for system, rolebinding_items in rolebinding_info.items():
        # check role - privilege mapping
        check_config__privilege_valid(system, rolebinding_items, privilege_info)
        # check role - org:member is only allow including path with {org_id}
        check_config__org_member_config(system, rolebinding_items, privilege_info)
        # check role - org:member is a basic role in org level
        check_config__basic_org_role_config(system, rolebinding_items, privilege_info)
        pass
    print("============ All configuration checks have passed.  ============")
    return
    
def notify_from_feishu(msg):
    import requests

    config_info = read_account_config()
    webhook_url = "https://open.feishu.cn/open-apis/bot/v2/hook/f904ccc5-f532-4ffb-91c1-c4ff54dbaa99"
    if webhook_url == "":
        print("get webhook_url failed. skip notify_compare_diff.")
        return
    print(webhook_url)
    headers = {"content-type": "application/json"}
    response = requests.request("POST", webhook_url, headers=headers, data=json.dumps(msg))
    print(response.text)
    return
    

def gen_multi_row_feishu_msg(title, line_list):
    content_line_list = [[{"tag":"text", "text": i}] for i in line_list]
    data_dict= {
            "msg_type": "post",
            "content": {
                "post": {
                    "en_us": {
                        "title": title,
                        "content": content_line_list
                        }
                    }
                }
            }
    return data_dict

def gen_rolebinding_diff_lines(delete_list, add_list):
    ret_lines = []
    # ret_lines.append("--------------------------------")
    # ret_lines.append("DELETE ROLEBINDINGS:")
    for item in delete_list:
        # line = " " + item["role"] + ": " + ", ".join(item["uri_list"])
        for uri in item["uri_list"]:
            line = "[Delete] uri: " + uri + " from role: " + item["role"]
            ret_lines.append(line)
            pass
        #line = " " + item["role"] + ": " + ", ".join(item["uri_list"])
        #ret_lines.append(line)
        pass
    #if len(delete_list) == 0:
    #    ret_lines.append(" No RoleBinding need to remove")
    #ret_lines.append("--------------------------------")
    #ret_lines.append("ADD ROLEBINDINGS:")
    for item in add_list:
        for uri in item["uri_list"]:
            line = "[Add] uri: " + uri + " to role: " + item["role"]
            ret_lines.append(line)
            pass
        #line = " " + item["role"] + ": " + ", ".join(item["uri_list"])
        #ret_lines.append(line)
        pass
    #if len(add_list) == 0:
    #    ret_lines.append(" No RoleBinding need to add")
    #ret_lines.append("--------------------------------")
    return ret_lines

def gen_compare_diff_feishu_msg(notification_dict):
    global env_param
    title = "["+env_param+"] Compare RoleBinding Diff Between DB and ConfigFile"
    if notification_dict == None or len(notification_dict) == 0:
        return gen_multi_row_feishu_msg(title, ['NO Diff Found'])
    line_list = []
    for system, role_dict in notification_dict.items():
        line_list.append("====== ["+system+"] =======")
        line_list.extend(gen_rolebinding_diff_lines(role_dict["delete_list"], role_dict["add_list"]))
        pass
    return gen_multi_row_feishu_msg(title, line_list)


def gen_update_diff_feishu_msg(system, delete_list, add_list):
    global env_param
    title = "["+env_param+"] (" + system + ") Update Rolebinding Configuration"
    line_list = []
    line_list.extend(gen_rolebinding_diff_lines(delete_list, add_list))
    return gen_multi_row_feishu_msg(title, line_list)

def compare_all_system_rolebindings(dump_rolebinding_file):
    system_list = get_system_list(dump_rolebinding_file)
    notification_dict = {}
    for system in system_list:
        print("+++++++++  compare system [" + system + "] ++++++++")
        delete_list, add_list = compare_config(system, dump_rolebinding_file)
        if len(delete_list) != 0 or len(add_list) !=0:
            notification_dict[system] = {"delete_list": delete_list, "add_list": add_list}
        pass
    if len(notification_dict) == 0:
        print("No system diff found, skip sending feishu message")
        return
    notify_msg = gen_compare_diff_feishu_msg(notification_dict)
    print(notify_msg)
    notify_from_feishu(notify_msg)
    return


# ---------- code for check next gen priv start ---------------
def get_dump_sys_privs_dict():
    data = read_json_from_file(dump_privilege_file)
    sys_privs_dict = {}
    for sys, data in data.items():
        for priv in data.keys():
            if sys not in sys_privs_dict:
                sys_privs_dict[sys] = []
            sys_privs_dict[sys].append(priv)
    return sys_privs_dict

def get_online_sys_privs_dict():
    ret_dict = {}
    connection = get_pymysql_connection()
    sql = "select resource_uri,`system` from rbac_privileges order by `system`,resource_uri"
    rows = get_sql(connection, sql)
    for row in rows:
        resource_uri = row[0]
        system = row[1]
        if system not in ret_dict:
            ret_dict[system] = []
            pass
        ret_dict[system].append(resource_uri)
    return ret_dict

def get_pymysql_connection():
    global global_pyconnection
    if global_pyconnection != None:
        return global_pyconnection
    mysql_conf = get_account_mysql_conf()
    global_pyconnection = pymysql.connect(
        host=mysql_conf['host'],
        port=mysql_conf['port'],
        user=mysql_conf['user'],
        password=mysql_conf['password'],
        database=mysql_conf['database']
        )
    return global_pyconnection

def get_account_mysql_conf():
    print("get pyconnection env: " + pyconnection_env)
    mysql_conf = {}
    if pyconnection_env == "dev":
        mysql_conf['host'] = 'nightly-cbp-dev-staging.cluster-c6owagykop0g.us-west-2.rds.amazonaws.com'
        mysql_conf['port'] = 3306
        mysql_conf['user'] = 'R9ez7VWZfnXH5Lw.nightly_account'
        mysql_conf['password'] = pyconnection_db_pass
        mysql_conf['database'] = 'nightly_ms_central'
    elif pyconnection_env == "staging":
        mysql_conf['host'] = 'nightly-cbp-dev-staging.cluster-c6owagykop0g.us-west-2.rds.amazonaws.com'
        mysql_conf['port'] = 3306
        mysql_conf['user'] = 'R9ez7VWZfnXH5Lw.staging'
        mysql_conf['password'] = pyconnection_db_pass
        mysql_conf['database'] = 'staging_central'
    return mysql_conf

def check_priv_valid(priv):
    if "clusters/{" in priv and "tidbs/{" not in priv:
        return False
    return True

def check_config__next_gen_privs_valid():
    dump_sys_privs_dict = get_dump_sys_privs_dict()
    online_sys_privs_dict = get_online_sys_privs_dict()
    for sys, privs in dump_sys_privs_dict.items():
        if sys not in online_sys_privs_dict:
            print(sys + " is not in online dict")
            for priv in privs:
                is_valid = check_priv_valid(priv)
                if not is_valid:
                    print(sys + ":" + priv  + " is invalid")
                    exit(1)
                    return
        else:
            online_privs = online_sys_privs_dict[sys]
            for priv in privs:
                if priv not in online_privs:
                    print(sys + ":" + priv  + " is not in online privs")
                    is_valid = check_priv_valid(priv)
                    if not is_valid:
                        print(sys + ":" + priv  + " is invalid")
                        exit(1)
                        return
# ---------- code for check next gen priv end ---------------

def process():
    if sys.argv[1] == "dump":
        dump_config()
    elif sys.argv[1] == "check-config":
        check_config()
    elif sys.argv[1] == "list-system":
        system_list = get_system_list(dump_rolebinding_file)
        print(system_list)
    elif sys.argv[1] == "compare" and len(sys.argv) > 2:
        compare_system = sys.argv[2]
        system_list = get_system_list(dump_rolebinding_file)
        if compare_system not in system_list:
            print("not found "+ compare_system + " in system_list")
            return 1
        compare_config(compare_system, dump_rolebinding_file)
    elif sys.argv[1] == "compare-all":
        compare_all_system_rolebindings(dump_rolebinding_file)
    elif sys.argv[1] == "update-all":
        system_list = get_system_list(dump_rolebinding_file)
        notification_dict = {}
        for system in system_list:
            delete_list, add_list = compare_config(system, dump_rolebinding_file)
            if len(delete_list) != 0 or len(add_list) != 0:
                update_config(system, dump_rolebinding_file, dump_privilege_file, False)
                notify_msg = gen_update_diff_feishu_msg(system, delete_list, add_list)
                print(notify_msg)
                notify_from_feishu(notify_msg)
                print("after update operation, post check the rbac diff between configuration and online")
                compare_config(system, dump_rolebinding_file)
                pass
            else:
                print("+++++++++  Skip update  system [" + system + "], No Diff found ++++++++")
            pass
        while True:
            compare_all_system_rolebindings(dump_rolebinding_file)
            time.sleep(86400)
        pass
    elif sys.argv[1] == "update" and len(sys.argv) > 2:
        compare_system = sys.argv[2]
        update_config(compare_system, dump_rolebinding_file, dump_privilege_file, False)
        print("after update operation, post check the rbac diff between configuration and online")
        compare_config(compare_system, dump_rolebinding_file)
    elif sys.argv[1] == "preview" and len(sys.argv) > 2:
        compare_system = sys.argv[2]
        update_config(compare_system, dump_rolebinding_file, dump_privilege_file, True)
    else:
        print("unsupport type")
        return 1
    return 0

if __name__ == "__main__":
    process()