1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
| import asyncio
import json
import os
import websockets
from typing import List, Optional
class CDPFileHandler:
"""CDP 文件处理器（上传 + 下载）"""
def __init__(self, ws, session_id=None):
self.ws = ws
self.session_id = session_id
self._cmd_id = 0
self._downloads = {}
async def _cdp(self, method: str, params: dict = None,
session_id: str = None) -> dict:
"""发送 CDP 命令"""
self._cmd_id += 1
msg = {"id": self._cmd_id, "method": method,
"params": params or {}}
sid = session_id or self.session_id
if sid:
msg["sessionId"] = sid
await self.ws.send(json.dumps(msg))
async for resp in self.ws:
data = json.loads(resp)
if data.get("id") == self._cmd_id:
return data.get("result", {})
async def enable_chooser_interception(self, enabled: bool = True):
"""启用/禁用文件选择拦截"""
await self._cdp("Page.setInterceptFileChooserDialog",
{"enabled": enabled})
async def handle_chooser(self, file_paths: List[str]):
"""处理文件选择对话框"""
await self._cdp("Page.handleFileChooser", {
"action": "accept",
"files": file_paths
})
async def upload_via_intercept(self, css_selector: str,
file_paths: List[str]) -> bool:
"""通过拦截文件对话框上传"""
abs_paths = [os.path.abspath(f) for f in file_paths
if os.path.exists(f)]
if not abs_paths:
raise FileNotFoundError("没有可用的文件")
await self.enable_chooser_interception(True)
future = asyncio.get_event_loop().create_future()
async def _wait_chooser():
async for msg in self.ws:
data = json.loads(msg)
if "id" in data:
continue
if data.get("method") == "Page.fileChooserOpened":
if not future.done():
future.set_result(True)
break
listener = asyncio.create_task(_wait_chooser())
await self._cdp("Runtime.evaluate", {
"expression": f"document.querySelector('{css_selector}').click()"
})
try:
await asyncio.wait_for(future, timeout=10)
await self.handle_chooser(abs_paths)
return True
except asyncio.TimeoutError:
return False
finally:
listener.cancel()
await self.enable_chooser_interception(False)
async def set_file_input(self, css_selector: str,
file_paths: List[str]):
"""直接给文件输入框设置文件"""
abs_paths = [os.path.abspath(f) for f in file_paths
if os.path.exists(f)]
result = await self._cdp("Runtime.evaluate", {
"expression": f"document.querySelector('{css_selector}')",
"objectGroup": "files"
})
object_id = result.get("result", {}).get("objectId")
if not object_id:
raise ValueError(f"元素未找到: {css_selector}")
await self._cdp("DOM.setFileInputFiles", {
"files": abs_paths,
"objectId": object_id
})
print(f"已设置文件: {[os.path.basename(f) for f in abs_paths]}")
async def set_download_path(self, path: str, behavior: str = "allow"):
"""设置下载路径（浏览器级命令）"""
os.makedirs(path, exist_ok=True)
await self._cdp("Browser.setDownloadBehavior", {
"behavior": behavior,
"downloadPath": os.path.abspath(path)
}, session_id=None)
print(f"下载路径: {os.path.abspath(path)}")
async def download_url(self, url: str,
save_dir: str = "./downloads") -> Optional[dict]:
"""下载一个 URL 并等待完成"""
save_dir = os.path.abspath(save_dir)
await self.set_download_path(save_dir)
result = await self._cdp("Target.createTarget", {"url": url},
session_id=None)
target_id = result.get("targetId")
if not target_id:
return None
await asyncio.sleep(3)
await self._cdp("Target.closeTarget", {"targetId": target_id},
session_id=None)
return {"targetId": target_id, "saveDir": save_dir}
async def listen_downloads(self, timeout: float = 30) -> List[dict]:
"""监听下载事件并返回结果"""
events = []
async def _handler(method: str, params: dict):
nonlocal events
if method == "Page.downloadWillBegin":
events.append({
"type": "begin",
"url": params.get("url", ""),
"filename": params.get("suggestedFilename", "")
})
print(f"下载开始: {params.get('suggestedFilename', '')}")
elif method == "Page.downloadProgress":
state = params.get("state", "")
received = params.get("receivedBytes", 0)
total = params.get("totalBytes", 0)
if state == "completed":
events.append({
"type": "completed",
"received": received,
"total": total
})
print(f"下载完成: {total} bytes")
elif state == "canceled":
events.append({"type": "canceled"})
print("下载取消")
try:
async with asyncio.timeout(timeout):
async for msg in self.ws:
data = json.loads(msg)
if "id" in data:
continue
await _handler(data.get("method", ""),
data.get("params", {}))
except asyncio.TimeoutError:
pass
return events
|
