335 lines
13 KiB
Python
335 lines
13 KiB
Python
import aiosqlite
|
||
import os
|
||
import asyncio
|
||
from core.log import logger
|
||
|
||
VALID_PROTOCOLS = ['http', 'https', 'socks4', 'socks5']
|
||
|
||
class SQLiteManager:
|
||
_instance = None
|
||
_connection = None
|
||
_lock = asyncio.Lock()
|
||
|
||
def __new__(cls, *args, **kwargs):
|
||
if cls._instance is None:
|
||
cls._instance = super(SQLiteManager, cls).__new__(cls)
|
||
return cls._instance
|
||
|
||
def __init__(self, db_path=None):
|
||
if hasattr(self, 'initialized') and self.initialized:
|
||
return
|
||
|
||
if db_path is None:
|
||
base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
db_dir = os.path.join(base_dir, 'db')
|
||
if not os.path.exists(db_dir):
|
||
os.makedirs(db_dir)
|
||
self.db_path = os.path.join(db_dir, 'proxies.sqlite')
|
||
else:
|
||
self.db_path = db_path
|
||
|
||
self.initialized = True
|
||
|
||
async def get_connection(self):
|
||
async with self._lock:
|
||
if self._connection is None:
|
||
self._connection = await aiosqlite.connect(self.db_path)
|
||
await self._connection.execute("PRAGMA journal_mode=WAL")
|
||
await self._connection.execute("PRAGMA synchronous=NORMAL")
|
||
await self._connection.execute("PRAGMA cache_size=-64000")
|
||
await self._connection.execute("PRAGMA temp_store=MEMORY")
|
||
return self._connection
|
||
|
||
async def close_connection(self):
|
||
async with self._lock:
|
||
if self._connection is not None:
|
||
await self._connection.close()
|
||
self._connection = None
|
||
|
||
async def init_db(self):
|
||
"""初始化数据库和表结构"""
|
||
db = await self.get_connection()
|
||
await db.execute('''
|
||
CREATE TABLE IF NOT EXISTS proxies (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
ip TEXT NOT NULL,
|
||
port INTEGER NOT NULL,
|
||
protocol TEXT DEFAULT 'http',
|
||
score INTEGER DEFAULT 10,
|
||
last_check TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||
UNIQUE(ip, port)
|
||
)
|
||
''')
|
||
|
||
await db.execute('CREATE INDEX IF NOT EXISTS idx_score ON proxies(score)')
|
||
await db.execute('CREATE INDEX IF NOT EXISTS idx_protocol ON proxies(protocol)')
|
||
await db.execute('CREATE INDEX IF NOT EXISTS idx_last_check ON proxies(last_check)')
|
||
await db.execute('CREATE INDEX IF NOT EXISTS idx_ip_port ON proxies(ip, port)')
|
||
|
||
await db.commit()
|
||
|
||
async def insert_proxy(self, ip, port, protocol='http', score=10):
|
||
"""异步插入或更新代理"""
|
||
try:
|
||
# 验证协议类型
|
||
if protocol not in VALID_PROTOCOLS:
|
||
protocol = 'http'
|
||
logger.warning(f"无效的协议类型 {protocol}，默认使用 http")
|
||
|
||
db = await self.get_connection()
|
||
# 先检查是否存在
|
||
async with db.execute('SELECT score FROM proxies WHERE ip = ? AND port = ?', (ip, port)) as cursor:
|
||
row = await cursor.fetchone()
|
||
if row:
|
||
# 如果存在，则更新最后检查时间和分数
|
||
await db.execute('''
|
||
UPDATE proxies SET last_check = CURRENT_TIMESTAMP, score = ?, protocol = ? WHERE ip = ? AND port = ?
|
||
''', (score, protocol, ip, port))
|
||
else:
|
||
# 如果不存在，则插入新记录
|
||
await db.execute('''
|
||
INSERT INTO proxies (ip, port, protocol, score, last_check)
|
||
VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP)
|
||
''', (ip, port, protocol, score))
|
||
await db.commit()
|
||
return True
|
||
except aiosqlite.IntegrityError as e:
|
||
# 处理唯一性约束冲突
|
||
if "UNIQUE" in str(e):
|
||
# 代理已存在，更新它
|
||
if protocol not in VALID_PROTOCOLS:
|
||
protocol = 'http'
|
||
db = await self.get_connection()
|
||
await db.execute('''
|
||
UPDATE proxies SET last_check = CURRENT_TIMESTAMP, score = ?, protocol = ? WHERE ip = ? AND port = ?
|
||
''', (score, protocol, ip, port))
|
||
await db.commit()
|
||
return True
|
||
else:
|
||
logger.error(f"数据库完整性错误: {e}")
|
||
return False
|
||
except Exception as e:
|
||
logger.error(f"插入代理失败 {ip}:{port} - {e}")
|
||
return False
|
||
|
||
async def get_all_proxies(self):
|
||
"""异步获取所有代理"""
|
||
db = await self.get_connection()
|
||
async with db.execute('SELECT ip, port, protocol, score, last_check FROM proxies') as cursor:
|
||
return await cursor.fetchall()
|
||
|
||
async def get_random_proxy(self):
|
||
"""异步随机获取一个高分代理"""
|
||
db = await self.get_connection()
|
||
async with db.execute('SELECT ip, port, protocol, score, last_check FROM proxies WHERE score > 0 ORDER BY RANDOM() LIMIT 1') as cursor:
|
||
return await cursor.fetchone()
|
||
|
||
async def update_score(self, ip, port, delta, min_score=0, max_score=100):
|
||
"""异步更新代理分数（增量更新，带分数限制）"""
|
||
try:
|
||
db = await self.get_connection()
|
||
# 获取当前分数
|
||
async with db.execute('SELECT score FROM proxies WHERE ip = ? AND port = ?', (ip, port)) as cursor:
|
||
row = await cursor.fetchone()
|
||
if row:
|
||
current_score = row[0]
|
||
new_score = max(min_score, min(max_score, current_score + delta))
|
||
await db.execute('''
|
||
UPDATE proxies SET score = ?, last_check = CURRENT_TIMESTAMP WHERE ip = ? AND port = ?
|
||
''', (new_score, ip, port))
|
||
if new_score <= 0:
|
||
await db.execute('DELETE FROM proxies WHERE score <= 0')
|
||
await db.commit()
|
||
return True
|
||
return False
|
||
except Exception as e:
|
||
logger.error(f"更新代理分数失败 {ip}:{port} - {e}")
|
||
return False
|
||
|
||
async def delete_proxy(self, ip, port):
|
||
"""异步删除指定代理"""
|
||
db = await self.get_connection()
|
||
await db.execute('DELETE FROM proxies WHERE ip = ? AND port = ?', (ip, port))
|
||
await db.commit()
|
||
|
||
async def count_proxies(self):
|
||
"""异步统计代理数量"""
|
||
db = await self.get_connection()
|
||
async with db.execute('SELECT COUNT(*) FROM proxies') as cursor:
|
||
row = await cursor.fetchone()
|
||
return row[0] if row else 0
|
||
|
||
async def get_proxies_paginated_with_total(self, page: int = 1, page_size: int = 20,
|
||
protocol: str = None, min_score: int = 0,
|
||
max_score: int = None,
|
||
sort_by: str = 'last_check',
|
||
sort_order: str = 'DESC'):
|
||
"""分页获取代理列表（一次查询返回数据和总数）"""
|
||
db = await self.get_connection()
|
||
conditions = ['score >= ?']
|
||
params = [min_score]
|
||
|
||
if protocol:
|
||
conditions.append('protocol = ?')
|
||
params.append(protocol)
|
||
|
||
if max_score is not None:
|
||
conditions.append('score <= ?')
|
||
params.append(max_score)
|
||
|
||
where_clause = ' AND '.join(conditions)
|
||
|
||
order_by_clause = f'{sort_by} {sort_order}'
|
||
|
||
offset = (page - 1) * page_size
|
||
query = f'''
|
||
SELECT ip, port, protocol, score, last_check,
|
||
COUNT(*) OVER() as total_count
|
||
FROM proxies
|
||
WHERE {where_clause}
|
||
ORDER BY {order_by_clause}
|
||
LIMIT ? OFFSET ?
|
||
'''
|
||
params.extend([page_size, offset])
|
||
|
||
async with db.execute(query, params) as cursor:
|
||
rows = await cursor.fetchall()
|
||
total = rows[0][5] if rows else 0
|
||
proxies = [(row[0], row[1], row[2], row[3], row[4]) for row in rows]
|
||
return proxies, total
|
||
|
||
async def get_proxies_paginated(self, page: int = 1, page_size: int = 20,
|
||
protocol: str = None, min_score: int = 0,
|
||
max_score: int = None,
|
||
sort_by: str = 'last_check',
|
||
sort_order: str = 'DESC'):
|
||
"""分页获取代理列表"""
|
||
db = await self.get_connection()
|
||
conditions = ['score >= ?']
|
||
params = [min_score]
|
||
|
||
if protocol:
|
||
conditions.append('protocol = ?')
|
||
params.append(protocol)
|
||
|
||
if max_score is not None:
|
||
conditions.append('score <= ?')
|
||
params.append(max_score)
|
||
|
||
where_clause = ' AND '.join(conditions)
|
||
|
||
order_by_clause = f'{sort_by} {sort_order}'
|
||
|
||
offset = (page - 1) * page_size
|
||
query = f'''
|
||
SELECT ip, port, protocol, score, last_check
|
||
FROM proxies
|
||
WHERE {where_clause}
|
||
ORDER BY {order_by_clause}
|
||
LIMIT ? OFFSET ?
|
||
'''
|
||
params.extend([page_size, offset])
|
||
|
||
async with db.execute(query, params) as cursor:
|
||
return await cursor.fetchall()
|
||
|
||
async def get_proxies_total(self, protocol: str = None, min_score: int = 0, max_score: int = None):
|
||
"""获取符合条件的代理总数"""
|
||
db = await self.get_connection()
|
||
conditions = ['score >= ?']
|
||
params = [min_score]
|
||
|
||
if protocol:
|
||
conditions.append('protocol = ?')
|
||
params.append(protocol)
|
||
|
||
if max_score is not None:
|
||
conditions.append('score <= ?')
|
||
params.append(max_score)
|
||
|
||
where_clause = ' AND '.join(conditions)
|
||
|
||
query = f'SELECT COUNT(*) FROM proxies WHERE {where_clause}'
|
||
|
||
async with db.execute(query, params) as cursor:
|
||
row = await cursor.fetchone()
|
||
return row[0] if row else 0
|
||
|
||
async def get_proxy_detail(self, ip: str, port: int):
|
||
"""获取单个代理的详细信息"""
|
||
db = await self.get_connection()
|
||
async with db.execute(
|
||
'SELECT ip, port, protocol, score, last_check FROM proxies WHERE ip = ? AND port = ?',
|
||
(ip, port)
|
||
) as cursor:
|
||
row = await cursor.fetchone()
|
||
return row
|
||
|
||
async def batch_delete_proxies(self, proxy_list: list):
|
||
"""批量删除代理，返回实际删除的数量"""
|
||
deleted_count = 0
|
||
db = await self.get_connection()
|
||
for ip, port in proxy_list:
|
||
cursor = await db.execute('DELETE FROM proxies WHERE ip = ? AND port = ?', (ip, port))
|
||
deleted_count += cursor.rowcount
|
||
await db.commit()
|
||
return deleted_count
|
||
|
||
async def get_stats(self):
|
||
"""获取统计信息"""
|
||
db = await self.get_connection()
|
||
stats = {}
|
||
|
||
async with db.execute('SELECT COUNT(*) FROM proxies') as cursor:
|
||
row = await cursor.fetchone()
|
||
stats['total'] = row[0] if row else 0
|
||
|
||
async with db.execute('SELECT COUNT(*) FROM proxies WHERE score > 0') as cursor:
|
||
row = await cursor.fetchone()
|
||
stats['available'] = row[0] if row else 0
|
||
|
||
async with db.execute('SELECT COUNT(*) FROM proxies WHERE protocol = "http"') as cursor:
|
||
row = await cursor.fetchone()
|
||
stats['http_count'] = row[0] if row else 0
|
||
|
||
async with db.execute('SELECT COUNT(*) FROM proxies WHERE protocol = "https"') as cursor:
|
||
row = await cursor.fetchone()
|
||
stats['https_count'] = row[0] if row else 0
|
||
|
||
async with db.execute('SELECT COUNT(*) FROM proxies WHERE protocol = "socks4"') as cursor:
|
||
row = await cursor.fetchone()
|
||
stats['socks4_count'] = row[0] if row else 0
|
||
|
||
async with db.execute('SELECT COUNT(*) FROM proxies WHERE protocol = "socks5"') as cursor:
|
||
row = await cursor.fetchone()
|
||
stats['socks5_count'] = row[0] if row else 0
|
||
|
||
async with db.execute('SELECT AVG(score) FROM proxies') as cursor:
|
||
row = await cursor.fetchone()
|
||
stats['avg_score'] = row[0] if row and row[0] else 0
|
||
|
||
return stats
|
||
|
||
async def get_today_new_count(self):
|
||
"""获取今日新增代理数量"""
|
||
try:
|
||
db = await self.get_connection()
|
||
query = '''
|
||
SELECT COUNT(*) FROM proxies
|
||
WHERE DATE(last_check) = DATE('now', 'localtime')
|
||
'''
|
||
async with db.execute(query) as cursor:
|
||
row = await cursor.fetchone()
|
||
return row[0] if row else 0
|
||
except Exception as e:
|
||
logger.error(f"获取今日新增数量失败: {e}")
|
||
return 0
|
||
|
||
async def clean_invalid_proxies(self):
|
||
"""清理无效代理（分数<=0）"""
|
||
db = await self.get_connection()
|
||
async with db.execute('DELETE FROM proxies WHERE score <= 0') as cursor:
|
||
deleted_count = cursor.rowcount
|
||
await db.commit()
|
||
return deleted_count
|