import requests
from bs4 import BeautifulSoup
from PIL import Image
import io

class AccountSpider:
    def __init__(self):
        self.base_url = "http://121.204.253.175:8088"
        self.session = requests.Session()
        # 设置默认请求头
        self.session.headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        }

    def get_verification_code(self,onlysave = False):
        """获取并保存验证码图片"""
        code_url = f"{self.base_url}/code.asp"
        response = self.session.get(code_url)
        
        # 保存验证码图片
        image = Image.open(io.BytesIO(response.content))
        image.save('/bot/danding-bot/danding_bot/plugins/damo_balance/verification_code.png')
        print("验证码图片已保存为 verification_code.png")
        # 仅保存验证码图片
        if onlysave:
            return
        # 等待用户输入验证码
        return input("请输入验证码: ")

    def login(self, username, password,v_code=""):
        """执行登录操作"""

        # 获取验证码
        if v_code:
            verification_code = v_code
        else:
            verification_code = self.get_verification_code()
        
        # 准备登录数据
        login_data = {
            'login_type': '0',
            'f_user': username,
            'f_code': password,
            'codeOK': verification_code,
            'Submit': '%C8%B7%B6%A8'
        }
        
        # 发送登录请求
        login_url = f"{self.base_url}/login_result.asp"
        response = self.session.post(login_url, data=login_data)
        response.encoding = 'gb2312'  # 设置正确的编码
        
        # 检查登录是否成功 - 通过检查是否包含重定向到account.asp的脚本
        if "window.location.href=\"account.asp\"" in response.text:
            return True
        return False

    def get_balance(self):
        """获取账户余额"""
        account_url = f"{self.base_url}/account.asp"
        response = self.session.get(account_url)
        response.encoding = 'gb2312'  # 设置正确的编码
        
        soup = BeautifulSoup(response.text, 'html.parser')
        balance_text = soup.find_all('span', class_='red')[1].text
        return float(balance_text)

def main():
    # 账号密码配置
    USERNAME = "xsllovemlj"
    PASSWORD = "xsl1314520mlj"
    
    spider = AccountSpider()
    
    # 尝试登录
    if spider.login(USERNAME, PASSWORD):
        print("登录成功！")
        balance = spider.get_balance()
        print(f"账户余额：{balance}元")
    else:
        print("登录失败，请检查账号密码或验证码是否正确")