su.py

import os
import json
import subprocess
import pytz
from datetime import datetime
from sucode import su

def check_password(username, entered_password):
    passwords = su(username)
    for service, password in passwords.items():
        if entered_password == password:
            return True, service
    return False, None

def find_user_by_user_id(user_id):
    data_path = "/mnt/upload/userdata.json"
    with open(data_path, 'r') as file:
        data = json.load(file)
    for user in data['users']:
        if user['userID'] == user_id:
            return user
    return None

def get_user_info():
    def aria2(url, filename, d):
        subprocess.run(
            ["aria2c", "-c", "-x", "16", "-s", "16", url, "-o", filename, "-d", d],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL
        )
    url = "https://modelscope.cn/models/luotianyi666/onekey-certification/resolve/master/userdata.pth"
    aria2(url, "userdata.json", "/mnt/upload")
    url = "https://modelscope.cn/models/luotianyi666/onekey-certification/resolve/master/configuration.json"
    aria2(url, "configuration.json", "/mnt/upload")

def check_repeat_username(username):
    user_info = find_user_by_user_id(username)
    if user_info:
        return True, None
    else:
        user_data = {
            "userID": user_info["userID"],
            "Aliyun_ID": user_info["Aliyun_ID"],
            "FirstTime": user_info["FirstTime"]
        }
        return False, user_data

def userdataupload():
    os.chdir("/mnt")
    os.rename("/mnt/upload/userdata.json", "/mnt/upload/userdata.pth")
    from modelscope.hub.api import HubApi
    api_key = '16cf6731-9441-4f98-b537-58ca240f1437'
    api = HubApi()
    api.login(api_key)
    api.push_model(
        model_id="luotianyi666/onekey-certification",
        model_dir="/mnt/upload"
    )

# 获取阿里云信息
def get_Aliyun_info():
    dsw_region = os.environ.get("dsw_region")
    accountID = os.environ.get("ALIYUN_ACCOUNT_ID")
    return dsw_region, accountID

# 获取当前实例区域信息
def checkregion():
    dsw_region, _ = get_Aliyun_info() 
    region_map = {
        "cn-beijing": "北京",
        "cn-shanghai": "上海",
        "cn-shenzhen": "深圳",
        "cn-hangzhou": "杭州",
    }
    dsw_region_cn = region_map.get(dsw_region, dsw_region)
    if dsw_region == "cn-beijing":
        print(f"当前{dsw_region_cn}区服务器资源拥挤，请选择北京区")
        return False
    return True

def find_user_by_Aliyun_id(Aliyun_id):
    with open("/mnt/upload/userdata.json", 'r') as file:
        data = json.load(file)
    for user in data['users']:
        if user['Aliyun_ID'] == Aliyun_id:
            return user
    return None

def webui():
    print("\033[42m webui success\033[0m\n\n")

def sd3():
    print("\033[42m sd3 success\033[0m\n\n")

def comfyui():
    print("\033[42m comfyui success\033[0m\n\n")

def forge():
    print("\033[42m forge success\033[0m\n\n")

def fooocus():
    print("\033[42m fooocus success\033[0m\n\n")

def facefusion():
    print("\033[42m facefusion success\033[0m\n\n")

def PaintsUNDO():
    print("\033[42m PaintsUNDO success\033[0m\n\n")

def IClight():
    print("\033[42m IClight success\033[0m\n\n")
    
def ChatTTS():
    print("\033[42m IClight success\033[0m\n\n")

def Omost():
    print("\033[42m Omost success\033[0m\n\n")
    
def main():
    if not checkregion():
        return
    get_user_info()
    accountID = get_Aliyun_info()[1]
    user = find_user_by_Aliyun_id(accountID)
    if user:
        print("\033[36m欢迎使用阿里云DSW一键部署服务\033[0m\n\n")
        services = [service[:-4] for service, flag in user.items() if flag and service.endswith("Flag")]
        options = services + ["other"]
        
        for i, option in enumerate(options):
            print(f"{i+1}. {option}")
        
        choice = int(input("请选择服务：")) - 1
        selected_service = options[choice]
        print(f"{selected_service} ")
        print(f"您的阿里云账户ID为{accountID}")
        if selected_service != "other":
            globals()[selected_service]()
        else:
            remaining_services = [service[:-4] for service, flag in user.items() if not flag and service.endswith("Flag")]
            for i, service in enumerate(remaining_services):
                print(f"{i+1}. {service}")
                
            choice = int(input("请选择服务：")) - 1
            selected_service = remaining_services[choice]
        
            username = user['userID']
            print(f"您阿里云账号绑定的用户名为{username}")
            entered_password = input(f"请输入{selected_service}服务的密码：")
        
            # 检查密码
            is_valid, service = check_password(username, entered_password)
            if is_valid and service == selected_service:
                print(f"密码正确，开始为您部署{selected_service}")
                globals()[selected_service]()
                user[f"{selected_service}Flag"] = True
                for s in remaining_services:
                    if s != selected_service:
                        user[f"{s}Flag"] = False
                data_path = "/mnt/upload/userdata.json"
                with open(data_path, 'r') as file:
                    data = json.load(file)
                for i, existing_user in enumerate(data['users']):
                    if existing_user['Aliyun_ID'] == user['Aliyun_ID']:
                        data['users'][i] = user
                        break
                with open(data_path, 'w') as file:
                    json.dump(data, file, indent=4)
                userdataupload()
            else:
                print("密码错误")
    else:
        print("欢迎使用阿里云DSW一键部署服务")        
        services = ["webui", "sd3", "comfyui", "forge", "fooocus", "facefusion", "PaintsUNDO", "IClight", "ChatTTS", "Omost"]
        for i, service in enumerate(services):
            print(f"{i+1}. {service}")
        
        choice = int(input("请输入需要的服务序号：")) - 1
        selected_service = services[choice]
        print(f"您的阿里云账户ID为{accountID}")
        username = input("请输入用户名（QQ号/微信号/咸鱼ID）：")
        entered_password = input(f"请输入 {selected_service} 服务的密码：")
        
        # 检查用户名是否与数据库中的某个 userID 匹配
        existing_user = find_user_by_user_id(username)
        if existing_user:
            print(f"{username}已被阿里云{existing_user['Aliyun_ID']}用户绑定，绑定时间为{existing_user['FirstTime']}")
            print("请重新运行并输入其他用户名")
        else:
            is_valid, service = check_password(username, entered_password)
            if is_valid and service == selected_service:
                print(f"密码正确，开始为您部署 {selected_service}")
                globals()[selected_service]()
                with open("/mnt/upload/userdata.json", 'r') as file:
                    data = json.load(file)
                new_user = {
                    "num": str(len(data['users']) + 1).zfill(3),
                    "userID": username,
                    "Aliyun_ID": accountID,
                    "FirstTime": datetime.now(pytz.timezone('Asia/Shanghai')).strftime('%Y%m%d%H%M'),
                    f"{selected_service}Flag": True
                }
                for s in services:
                    if s != selected_service:
                        new_user[f"{s}Flag"] = False
                data['users'].append(new_user)
                with open("/mnt/upload/userdata.json", 'w') as file:
                    json.dump(data, file, indent=4)
                userdataupload()
            else:
                print("密码错误")
if __name__ == "__main__":
    main()