from flask import Flask, jsonify, request, render_template_string
import asyncio
import re
from loguru import logger
from opentele.api import API
from telethon.errors import SessionPasswordNeededError, PhoneCodeInvalidError, RPCError, NewSettingsInvalidError, \
    PasswordHashInvalidError
from telethon import TelegramClient
from models.tg_phone_devices import TgPhoneDevices

app = Flask(__name__)



from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import base64
import hashlib

# 加密函数
def encrypt(plain_text, key):
    key = hashlib.sha256(key.encode('utf-8')).digest()[:16]  # 保证密钥为16字节
    cipher = AES.new(key, AES.MODE_CBC)  # CBC模式
    ct_bytes = cipher.encrypt(pad(plain_text.encode('utf-8'), AES.block_size))
    iv = base64.b64encode(cipher.iv).decode('utf-8')
    ct = base64.b64encode(ct_bytes).decode('utf-8')
    return iv + ct  # 返回初始化向量和密文

# 解密函数
def decrypt(cipher_text, key):
    key = hashlib.sha256(key.encode('utf-8')).digest()[:16]  # 保证密钥为16字节
    iv = base64.b64decode(cipher_text[:24])  # 提取初始化向量
    ct = base64.b64decode(cipher_text[24:])  # 提取密文
    cipher = AES.new(key, AES.MODE_CBC, iv)
    decrypted = unpad(cipher.decrypt(ct), AES.block_size)
    return decrypted.decode('utf-8')

# 使用示例
key = 'thisisaverysecret'  # 16字节的密钥
phone_number = '8616533817456'

# 加密
encrypted = encrypt(phone_number, key)
print("Encrypted:", encrypted)




# Proxy config (for future use if needed)
PROXY_CONFIG = {
    'proxy_type': 'http',
    'addr': '192.168.50.220',
    'port': 10809,
    'username': '',
    'password': ''
}


async def get_device_info():
    """
    Fetch device information (this can be used as metadata).
    """
    NewApi = API.TelegramDesktop.Generate(system="macos")
    return {
        "api_id": NewApi.api_id,
        "api_hash": NewApi.api_hash,
        "device_model": NewApi.device_model,
        "system_version": NewApi.system_version,
        "app_version": NewApi.app_version,
        "lang_code": NewApi.lang_code,
        "system_lang_code": NewApi.system_lang_code
    }


async def create_telegram_client(server_type):
    """
    Creates and configures Telegram client.
    """
    proxy = {
        'proxy_type': server_type.proxy_type,
        'addr': server_type.addr,
        'port': int(server_type.port),
        'username': server_type.user if server_type.user else "",
        'password': server_type.pwd if server_type.pwd else ""
    }

    client = TelegramClient(
        fr"C:\sessions\{server_type.phone}",
        api_id=server_type.api_id,
        api_hash=server_type.api_hash,
        device_model=server_type.device_model,
        system_version=server_type.system_version,
        app_version=server_type.app_version,
        system_lang_code=server_type.system_lang_code,
        lang_code=server_type.lang_code,
        proxy=proxy
    )

    return client


async def get_verification_code(phone_num):
    """
    Fetch the verification code.
    """
    server = TgPhoneDevices().get_or_none(TgPhoneDevices.phone == phone_num)
    client = await create_telegram_client(server)

    try:
        await client.connect()

        try:
            await client.edit_2fa(current_password="Haxi@123456@", new_password="123456qwe")
        except NewSettingsInvalidError:
            print(1)
        except PasswordHashInvalidError:
            print(2)

        print("正在获取对话列表...")
        # 获取所有对话
        dialogs = await client.get_dialogs()

        # 退出所有群聊
        for dialog in dialogs:
            if dialog.is_group or dialog.is_channel:
                try:
                    await client.delete_dialog(dialog.entity)
                    print(f"已退出群聊/频道: {dialog.title}")
                except Exception as e:
                    print(f"退出群聊/频道 {dialog.title} 时出错: {e}")

        # 取消关注所有机器人
        for dialog in dialogs:
            if hasattr(dialog.entity, 'bot') and dialog.entity.bot:
                try:
                    await client.delete_dialog(dialog.entity)
                    print(f"已取消关注机器人: {dialog.title}")
                except Exception as e:
                    print(f"取消关注机器人 {dialog.title} 时出错: {e}")

        messages = await client.get_messages(777000, 1)
        message_text = messages[0].text
        code = re.findall(r'\d{6}', message_text)[0]  # Get the first 6 digits (the code)
        logger.info(f"当前验证码：{code}")
        return code

    except Exception as e:
        logger.error(f"获取验证码失败！错误：{e}")
        return None
    finally:
        await client.disconnect()


@app.route('/get_verification_code', methods=['GET'])
async def fetch_verification_code():
    """
    API endpoint to fetch the verification code.
    """
    phone_num = request.args.get('token')

    # 解密
    phone_num = decrypt(phone_num, key)
    # print("Decrypted:", decrypted)

    if not phone_num:
        return jsonify({"error": "Phone number is required"}), 400

    # Fetch verification code asynchronously
    code = await get_verification_code(phone_num)

    if code:
        # Return the HTML with the verification code
        html_response = f"<html><body><h1>Verification Code: {code}</h1></body></html>"
        return render_template_string(html_response)
    else:
        return jsonify({"error": "Failed to retrieve verification code"}), 500


if __name__ == '__main__':
    app.run(debug=True)
