198 lines
6.7 KiB
Python
198 lines
6.7 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
|
||
import time
|
||
import logging
|
||
import random
|
||
import string
|
||
import json
|
||
import concurrent.futures
|
||
import argparse
|
||
from email_api_client import EmailApiClient
|
||
|
||
# 配置日志
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||
)
|
||
logger = logging.getLogger("BatchRegistration")
|
||
|
||
def generate_random_username(length=8):
|
||
"""生成随机用户名"""
|
||
letters = string.ascii_lowercase + string.digits
|
||
return ''.join(random.choice(letters) for _ in range(length))
|
||
|
||
def register_account(client, domain, username=None):
|
||
"""
|
||
模拟注册账号流程
|
||
|
||
参数:
|
||
client: EmailApiClient实例
|
||
domain: 邮箱域名
|
||
username: 指定用户名,不指定则随机生成
|
||
|
||
返回:
|
||
注册结果字典
|
||
"""
|
||
# 生成或使用提供的用户名
|
||
username = username or generate_random_username()
|
||
email_address = f"{username}@{domain}"
|
||
|
||
logger.info(f"开始注册账号: {email_address}")
|
||
|
||
# 模拟注册操作
|
||
logger.info(f"模拟向注册服务发送注册请求...")
|
||
# 这里是模拟,实际中应该调用真实的注册API
|
||
time.sleep(1) # 模拟API调用延迟
|
||
|
||
# 等待验证邮件
|
||
logger.info(f"等待验证邮件...")
|
||
start_time = time.time()
|
||
email = client.wait_for_email(
|
||
email_address,
|
||
timeout=30, # 30秒超时
|
||
subject="验证", # 查找主题包含"验证"的邮件
|
||
keyword="验证码" # 内容包含"验证码"的邮件
|
||
)
|
||
|
||
wait_time = time.time() - start_time
|
||
|
||
if email:
|
||
# 提取验证码
|
||
verification_code = client.extract_verification_code(email)
|
||
|
||
# 模拟验证过程
|
||
if verification_code:
|
||
logger.info(f"提取到验证码: {verification_code},正在验证...")
|
||
# 模拟提交验证码
|
||
time.sleep(0.5) # 模拟API调用延迟
|
||
|
||
return {
|
||
'success': True,
|
||
'email_address': email_address,
|
||
'verification_code': verification_code,
|
||
'wait_time': wait_time,
|
||
'message': '注册成功'
|
||
}
|
||
else:
|
||
logger.warning(f"未能提取验证码")
|
||
return {
|
||
'success': False,
|
||
'email_address': email_address,
|
||
'wait_time': wait_time,
|
||
'message': '未能提取验证码'
|
||
}
|
||
else:
|
||
logger.warning(f"未收到验证邮件")
|
||
return {
|
||
'success': False,
|
||
'email_address': email_address,
|
||
'wait_time': wait_time,
|
||
'message': '未收到验证邮件'
|
||
}
|
||
|
||
def batch_register(domain, count, concurrent=2, api_url=None):
|
||
"""
|
||
批量注册账号
|
||
|
||
参数:
|
||
domain: 邮箱域名
|
||
count: 注册数量
|
||
concurrent: 并发数
|
||
api_url: API基础URL,默认为None使用客户端默认值
|
||
|
||
返回:
|
||
注册结果列表
|
||
"""
|
||
logger.info(f"开始批量注册 {count} 个账号,域名: {domain},并发数: {concurrent}")
|
||
|
||
# 创建API客户端
|
||
client = EmailApiClient(api_base_url=api_url) if api_url else EmailApiClient()
|
||
|
||
# 检查系统状态
|
||
status = client.check_system_status()
|
||
if not status.get('success', False):
|
||
logger.error(f"邮件系统状态异常: {status}")
|
||
return []
|
||
|
||
# 生成用户名列表
|
||
usernames = [generate_random_username() for _ in range(count)]
|
||
|
||
results = []
|
||
|
||
# 使用线程池并发注册
|
||
if concurrent > 1:
|
||
with concurrent.futures.ThreadPoolExecutor(max_workers=concurrent) as executor:
|
||
# 提交任务
|
||
future_to_username = {
|
||
executor.submit(register_account, client, domain, username): username
|
||
for username in usernames
|
||
}
|
||
|
||
# 获取结果
|
||
for future in concurrent.futures.as_completed(future_to_username):
|
||
username = future_to_username[future]
|
||
try:
|
||
result = future.result()
|
||
results.append(result)
|
||
logger.info(f"账号 {username}@{domain} 注册完成: {result['success']}")
|
||
except Exception as e:
|
||
logger.error(f"账号 {username}@{domain} 注册出错: {str(e)}")
|
||
results.append({
|
||
'success': False,
|
||
'email_address': f"{username}@{domain}",
|
||
'message': f'注册过程出错: {str(e)}'
|
||
})
|
||
else:
|
||
# 串行注册
|
||
for username in usernames:
|
||
try:
|
||
result = register_account(client, domain, username)
|
||
results.append(result)
|
||
logger.info(f"账号 {username}@{domain} 注册完成: {result['success']}")
|
||
except Exception as e:
|
||
logger.error(f"账号 {username}@{domain} 注册出错: {str(e)}")
|
||
results.append({
|
||
'success': False,
|
||
'email_address': f"{username}@{domain}",
|
||
'message': f'注册过程出错: {str(e)}'
|
||
})
|
||
|
||
# 统计结果
|
||
success_count = sum(1 for r in results if r.get('success', False))
|
||
fail_count = len(results) - success_count
|
||
|
||
logger.info(f"批量注册完成: 成功 {success_count} 个, 失败 {fail_count} 个")
|
||
|
||
return results
|
||
|
||
def save_results(results, filename):
|
||
"""将结果保存到文件"""
|
||
with open(filename, 'w', encoding='utf-8') as f:
|
||
json.dump(results, f, ensure_ascii=False, indent=2)
|
||
logger.info(f"结果已保存到文件: {filename}")
|
||
|
||
def main():
|
||
"""主函数"""
|
||
parser = argparse.ArgumentParser(description='批量注册账号示例')
|
||
parser.add_argument('--domain', type=str, default='nosqli.com', help='邮箱域名')
|
||
parser.add_argument('--count', type=int, default=5, help='注册数量')
|
||
parser.add_argument('--concurrent', type=int, default=2, help='并发数')
|
||
parser.add_argument('--output', type=str, default='registration_results.json', help='结果输出文件')
|
||
parser.add_argument('--api-url', type=str, help='邮件系统API基础URL,默认为http://74.48.75.19:5000/api')
|
||
|
||
args = parser.parse_args()
|
||
|
||
# 执行批量注册
|
||
results = batch_register(args.domain, args.count, args.concurrent, args.api_url)
|
||
|
||
# 保存结果
|
||
if results:
|
||
save_results(results, args.output)
|
||
|
||
# 退出代码
|
||
success_count = sum(1 for r in results if r.get('success', False))
|
||
return 0 if success_count == args.count else 1
|
||
|
||
if __name__ == "__main__":
|
||
exit(main()) |