network.py (9771B)
1 from enum import Enum 2 from typing import Any, Dict, List, Mapping, MutableMapping, Optional, Union 3 4 from ._module import BidiModule, command 5 from ..undefined import UNDEFINED, Maybe 6 7 8 class AuthCredentials(Dict[str, Any]): 9 def __init__(self, username: str, password: str): 10 dict.__init__(self, type="password", username=username, password=password) 11 12 13 class CacheBehavior(Enum): 14 BYPASS = "bypass" 15 DEFAULT = "default" 16 17 18 class NetworkBase64Value(Dict[str, Any]): 19 def __init__(self, value: str): 20 dict.__init__(self, type="base64", value=value) 21 22 23 class NetworkStringValue(Dict[str, Any]): 24 def __init__(self, value: str): 25 dict.__init__(self, type="string", value=value) 26 27 28 NetworkBytesValue = Union[NetworkStringValue, NetworkBase64Value] 29 30 31 class CookieHeader(Dict[str, Any]): 32 def __init__(self, name: str, value: NetworkBytesValue): 33 dict.__init__(self, name=name, value=value) 34 35 36 class Header(Dict[str, Any]): 37 def __init__(self, name: str, value: NetworkBytesValue): 38 dict.__init__(self, name=name, value=value) 39 40 41 class SetCookieHeader(Dict[str, Any]): 42 def __init__( 43 self, 44 name: str, 45 value: NetworkBytesValue, 46 domain: Optional[str] = None, 47 expiry: Optional[str] = None, 48 http_only: Optional[bool] = None, 49 max_age: Optional[int] = None, 50 path: Optional[str] = None, 51 same_site: Optional[str] = None, 52 secure: Optional[bool] = None, 53 ): 54 dict.__init__(self, name=name, value=value) 55 56 if domain is not None: 57 self["domain"] = domain 58 59 if expiry is not None: 60 self["expiry"] = expiry 61 62 if http_only is not None: 63 self["httpOnly"] = http_only 64 65 if max_age is not None: 66 self["maxAge"] = max_age 67 68 if path is not None: 69 self["path"] = path 70 71 if same_site is not None: 72 self["sameSite"] = same_site 73 74 if secure is not None: 75 self["secure"] = secure 76 77 78 class URLPatternPattern(Dict[str, Any]): 79 def __init__( 80 self, 81 protocol: Optional[str] = None, 82 hostname: Optional[str] = None, 83 port: Optional[str] = None, 84 pathname: Optional[str] = None, 85 search: Optional[str] = None, 86 ): 87 dict.__init__(self, type="pattern") 88 89 if protocol is not None: 90 self["protocol"] = protocol 91 92 if hostname is not None: 93 self["hostname"] = hostname 94 95 if port is not None: 96 self["port"] = port 97 98 if pathname is not None: 99 self["pathname"] = pathname 100 101 if search is not None: 102 self["search"] = search 103 104 105 class URLPatternString(Dict[str, Any]): 106 def __init__(self, pattern: str): 107 dict.__init__(self, type="string", pattern=pattern) 108 109 110 URLPattern = Union[URLPatternPattern, URLPatternString] 111 112 113 class Network(BidiModule): 114 @command 115 def add_intercept( 116 self, phases: List[str], url_patterns: Optional[List[URLPattern]] = None, contexts: Optional[List[str]] = None 117 ) -> Mapping[str, Any]: 118 params: MutableMapping[str, Any] = { 119 "phases": phases, 120 } 121 122 if url_patterns is not None: 123 params["urlPatterns"] = url_patterns 124 125 if contexts is not None: 126 params["contexts"] = contexts 127 128 return params 129 130 @add_intercept.result 131 def _add_intercept(self, result: Mapping[str, Any]) -> Any: 132 assert result["intercept"] is not None 133 return result["intercept"] 134 135 @command 136 def continue_with_auth( 137 self, 138 request: str, 139 action: str, 140 credentials: Optional[AuthCredentials] = None 141 ) -> Mapping[str, Any]: 142 params: MutableMapping[str, Any] = { 143 "request": request, 144 "action": action, 145 } 146 147 if action == "provideCredentials" and credentials is not None: 148 params["credentials"] = credentials 149 150 return params 151 152 @command 153 def continue_request(self, 154 request: str, 155 body: Optional[NetworkBytesValue] = None, 156 cookies: Optional[List[CookieHeader]] = None, 157 headers: Optional[List[Header]] = None, 158 method: Optional[str] = None, 159 url: Optional[str] = None) -> Mapping[str, Any]: 160 params: MutableMapping[str, Any] = { 161 "request": request, 162 } 163 164 if body is not None: 165 params["body"] = body 166 167 if cookies is not None: 168 params["cookies"] = cookies 169 170 if headers is not None: 171 params["headers"] = headers 172 173 if method is not None: 174 params["method"] = method 175 176 if url is not None: 177 params["url"] = url 178 179 return params 180 181 @command 182 def continue_response( 183 self, 184 request: str, 185 cookies: Optional[List[SetCookieHeader]] = None, 186 credentials: Optional[AuthCredentials] = None, 187 headers: Optional[List[Header]] = None, 188 reason_phrase: Optional[str] = None, 189 status_code: Optional[int] = None) -> Mapping[str, Any]: 190 params: MutableMapping[str, Any] = { 191 "request": request, 192 } 193 194 if cookies is not None: 195 params["cookies"] = cookies 196 197 if credentials is not None: 198 params["credentials"] = credentials 199 200 if headers is not None: 201 params["headers"] = headers 202 203 if reason_phrase is not None: 204 params["reasonPhrase"] = reason_phrase 205 206 if status_code is not None: 207 params["statusCode"] = status_code 208 209 210 return params 211 212 @command 213 def fail_request(self, request: str) -> Mapping[str, Any]: 214 params: MutableMapping[str, Any] = {"request": request} 215 return params 216 217 @command 218 def provide_response( 219 self, 220 request: str, 221 body: Optional[NetworkBytesValue] = None, 222 cookies: Optional[List[SetCookieHeader]] = None, 223 headers: Optional[List[Header]] = None, 224 reason_phrase: Optional[str] = None, 225 status_code: Optional[int] = None) -> Mapping[str, Any]: 226 params: MutableMapping[str, Any] = { 227 "request": request, 228 } 229 230 if body is not None: 231 params["body"] = body 232 233 if cookies is not None: 234 params["cookies"] = cookies 235 236 if headers is not None: 237 params["headers"] = headers 238 239 if reason_phrase is not None: 240 params["reasonPhrase"] = reason_phrase 241 242 if status_code is not None: 243 params["statusCode"] = status_code 244 245 return params 246 247 @command 248 def remove_intercept(self, intercept: str) -> Mapping[str, Any]: 249 params: MutableMapping[str, Any] = {"intercept": intercept} 250 return params 251 252 @command 253 def add_data_collector( 254 self, 255 data_types: List[str], 256 max_encoded_data_size: int, 257 collector_type: Optional[str] = None, 258 contexts: Optional[List[str]] = None, 259 user_contexts: Optional[List[str]] = None) -> Mapping[str, Any]: 260 params: MutableMapping[str, Any] = { 261 "dataTypes": data_types, 262 "maxEncodedDataSize": max_encoded_data_size, 263 } 264 265 if collector_type is not None: 266 params["collectorType"] = collector_type 267 268 if contexts is not None: 269 params["contexts"] = contexts 270 271 if user_contexts is not None: 272 params["userContexts"] = user_contexts 273 274 return params 275 276 @add_data_collector.result 277 def _add_data_collector(self, result: Mapping[str, Any]) -> Any: 278 assert result["collector"] is not None 279 return result["collector"] 280 281 @command 282 def remove_data_collector(self, collector: str) -> Mapping[str, Any]: 283 params: MutableMapping[str, Any] = {"collector": collector} 284 return params 285 286 @command 287 def disown_data( 288 self, 289 request: str, 290 data_type: str, 291 collector: str) -> Mapping[str, Any]: 292 params: MutableMapping[str, Any] = { 293 "request": request, 294 "dataType": data_type, 295 "collector": collector, 296 } 297 return params 298 299 @command 300 def set_cache_behavior( 301 self, 302 cache_behavior: CacheBehavior, 303 contexts: Optional[List[str]] = None) -> Mapping[str, Any]: 304 params: MutableMapping[str, Any] = {"cacheBehavior": cache_behavior} 305 306 if contexts is not None: 307 params["contexts"] = contexts 308 309 return params 310 311 @command 312 def get_data( 313 self, 314 request: str, 315 data_type: str, 316 collector: Optional[str] = None, 317 disown: Optional[bool] = None) -> Mapping[str, Any]: 318 params: MutableMapping[str, Any] = { 319 "request": request, 320 "dataType": data_type, 321 } 322 323 if collector is not None: 324 params["collector"] = collector 325 326 if disown is not None: 327 params["disown"] = disown 328 329 return params 330 331 @get_data.result 332 def _get_data(self, result: Mapping[str, Any]) -> Any: 333 assert result["bytes"] is not None 334 return result["bytes"] 335 336 @command 337 def set_extra_headers( 338 self, 339 headers: List[Dict[str, Any]], 340 contexts: Maybe[List[str]] = UNDEFINED, 341 user_contexts: Maybe[List[str]] = UNDEFINED, 342 ) -> Mapping[str, Any]: 343 return { 344 "headers": headers, 345 "contexts": contexts, 346 "userContexts": user_contexts, 347 }