zhitou_trade/connect_wifi.py

import os
import time
import subprocess
import requests
import ctypes
import sys
from typing import Tuple

import logging
logger = logging.getLogger(__name__)

def is_admin() -> bool:
    """检查管理员权限"""
    try:
        return ctypes.windll.shell32.IsUserAnAdmin()
    except:
        return False

def check_internet_connection(test_urls=None, timeout=5) -> bool:
    """检测网络连接"""
    if test_urls is None:
        test_urls = [
            "https://www.baidu.com",
            "https://www.qq.com",
            "https://www.cloudflare.com"
        ]

    for url in test_urls:
        try:
            response = requests.get(url, timeout=timeout)
            if response.status_code == 200:
                return True
        except:
            continue
    return False

def run_netsh_command(command: list, timeout: int = 10) -> Tuple[bool, str]:
    """安全执行命令"""
    try:
        result = subprocess.run(
            command,
            check=True,
            shell=True,
            timeout=timeout,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            universal_newlines=False  # 捕获字节输出
        )
        # 解码输出，忽略错误
        stdout = result.stdout.decode('utf-8', errors='ignore').strip()
        stderr = result.stderr.decode('utf-8', errors='ignore').strip()
        output = stdout if stdout else stderr
        return True, output
    except subprocess.CalledProcessError as e:
        # 解码异常中的输出
        stdout = e.stdout.decode('utf-8', errors='ignore').strip()
        stderr = e.stderr.decode('utf-8', errors='ignore').strip()
        output = stderr if stderr else stdout or "Unknown error"
        return False, f"{output} (exit code {e.returncode})"
    except subprocess.TimeoutExpired as e:
        # 处理超时输出
        stdout = e.stdout.decode('utf-8', errors='ignore').strip() if e.stdout else ""
        stderr = e.stderr.decode('utf-8', errors='ignore').strip() if e.stderr else ""
        output = f"{stdout} {stderr}".strip()
        return False, f"Timeout after {timeout}s: {output}"
    except Exception as e:
        return False, f"Unexpected error: {str(e)}"

def connect_wifi(ssid: str = 'MaxEntropy', password: str = 'cskj12345678') -> bool:
    # 权限检查
    if not is_admin():
        print("请以管理员身份运行！")
        return False

    # 网络检测
    if check_internet_connection():
        print("当前网络正常")
        return True

    # 生成配置文件
    profile_path = f"wifi_temp_{ssid}.xml"
    try:
        # 生成配置文件内容
        wifi_config = f"""<?xml version="1.0"?>
<WLANProfile xmlns="http://www.microsoft.com/networking/WLAN/profile/v1">
    <name>{ssid}</name>
    <SSIDConfig>
        <SSID><name>{ssid}</name></SSID>
    </SSIDConfig>
    <connectionType>ESS</connectionType>
    <connectionMode>auto</connectionMode>
    <MSM>
        <security>
            <authEncryption>
                <authentication>WPA2PSK</authentication>
                <encryption>AES</encryption>
                <useOneX>false</useOneX>
            </authEncryption>
            <sharedKey>
                <keyType>passPhrase</keyType>
                <protected>false</protected>
                <keyMaterial>{password}</keyMaterial>
            </sharedKey>
        </security>
    </MSM>
</WLANProfile>"""

        # 写入临时文件（使用UTF-8编码）
        with open(profile_path, "w", encoding='utf-8') as f:
            f.write(wifi_config.strip())

        # 删除旧配置文件（如果有）
        run_netsh_command(["netsh", "wlan", "delete", "profile", ssid])

        # 添加配置文件
        success, msg = run_netsh_command([
            "netsh", "wlan", "add", "profile",
            f"filename={profile_path}"
        ])
        if not success:
            print(f"[添加配置文件失败] {msg}")
            return False

        # 执行连接
        success, msg = run_netsh_command([
            "netsh", "wlan", "connect",
            f"name={ssid}",
            "interface=Wi-Fi"
        ])
        if not success:
            print(f"[连接失败] {msg}")
            return False

        # 等待连接
        for i in range(5):
            time.sleep(2)
            if check_internet_connection():
                print(f"成功连接至 {ssid}")
                return True
            print(f"等待连接中...({i + 1}/5)")

        print("连接超时")
        return False

    finally:
        # 清理临时文件
        if os.path.exists(profile_path):
            try:
                os.remove(profile_path)
            except Exception as e:
                print(f"清理临时文件失败: {str(e)}")

if __name__ == "__main__":
    connect_wifi(ssid="ZTE_A5DA76", password="1234567890")