1717import base64
1818import os
1919import socket
20+ from enum import Enum
2021from typing import Optional
2122from urllib import parse
2223
2627from selenium .webdriver .common .proxy import ProxyType
2728
2829
30+ class AuthType (Enum ):
31+ BASIC = "Basic"
32+ BEARER = "Bearer"
33+ X_API_KEY = "X-API-Key"
34+
35+
2936class ClientConfig :
3037 def __init__ (
3138 self ,
@@ -38,8 +45,10 @@ def __init__(
3845 ca_certs : Optional [str ] = None ,
3946 username : Optional [str ] = None ,
4047 password : Optional [str ] = None ,
41- auth_type : Optional [str ] = "Basic" ,
48+ auth_type : Optional [AuthType ] = AuthType . BASIC ,
4249 token : Optional [str ] = None ,
50+ user_agent : Optional [str ] = None ,
51+ extra_headers : Optional [dict ] = None ,
4352 ) -> None :
4453 self .remote_server_addr = remote_server_addr
4554 self .keep_alive = keep_alive
@@ -51,6 +60,8 @@ def __init__(
5160 self .password = password
5261 self .auth_type = auth_type
5362 self .token = token
63+ self .user_agent = user_agent
64+ self .extra_headers = extra_headers
5465
5566 self .timeout = (
5667 (
@@ -198,14 +209,17 @@ def password(self, value: str) -> None:
198209 self ._password = value
199210
200211 @property
201- def auth_type (self ) -> str :
212+ def auth_type (self ) -> AuthType :
202213 """Returns the type of authentication to the remote server."""
203214 return self ._auth_type
204215
205216 @auth_type .setter
206- def auth_type (self , value : str ) -> None :
217+ def auth_type (self , value : AuthType ) -> None :
207218 """Sets the type of authentication to the remote server if it is not
208- using basic with username and password."""
219+ using basic with username and password.
220+
221+ :Args: value - AuthType enum value. For others, please use `extra_headers` instead
222+ """
209223 self ._auth_type = value
210224
211225 @property
@@ -219,6 +233,26 @@ def token(self, value: str) -> None:
219233 auth_type is not basic."""
220234 self ._token = value
221235
236+ @property
237+ def user_agent (self ) -> str :
238+ """Returns user agent to be added to the request headers."""
239+ return self ._user_agent
240+
241+ @user_agent .setter
242+ def user_agent (self , value : str ) -> None :
243+ """Sets user agent to be added to the request headers."""
244+ self ._user_agent = value
245+
246+ @property
247+ def extra_headers (self ) -> dict :
248+ """Returns extra headers to be added to the request."""
249+ return self ._extra_headers
250+
251+ @extra_headers .setter
252+ def extra_headers (self , value : dict ) -> None :
253+ """Sets extra headers to be added to the request."""
254+ self ._extra_headers = value
255+
222256 def get_proxy_url (self ) -> Optional [str ]:
223257 """Returns the proxy URL to use for the connection."""
224258 proxy_type = self .proxy .proxy_type
@@ -246,13 +280,12 @@ def get_proxy_url(self) -> Optional[str]:
246280
247281 def get_auth_header (self ) -> Optional [dict ]:
248282 """Returns the authorization to add to the request headers."""
249- auth_type = self .auth_type .lower ()
250- if auth_type == "basic" and self .username and self .password :
283+ if self .auth_type is AuthType .BASIC and self .username and self .password :
251284 credentials = f"{ self .username } :{ self .password } "
252285 encoded_credentials = base64 .b64encode (credentials .encode ("utf-8" )).decode ("utf-8" )
253- return {"Authorization" : f"Basic { encoded_credentials } " }
254- if auth_type == "bearer" and self .token :
255- return {"Authorization" : f"Bearer { self .token } " }
256- if auth_type == "oauth" and self .token :
257- return {"Authorization " : f"OAuth { self .token } " }
286+ return {"Authorization" : f"{ AuthType . BASIC . value } { encoded_credentials } " }
287+ if self . auth_type is AuthType . BEARER and self .token :
288+ return {"Authorization" : f"{ AuthType . BEARER . value } { self .token } " }
289+ if self . auth_type is AuthType . X_API_KEY and self .token :
290+ return {f" { AuthType . X_API_KEY . value } " : f"{ self .token } " }
258291 return None
0 commit comments