11from dataclasses import dataclass
22
33import json
4+ import re
45from base64 import b64encode , b64decode
56
67APERTURE_CLOUD = ".cloud.aperturedata.io"
@@ -24,10 +25,13 @@ class Configuration:
2425 # Max number of attempts to retry the initial connection (0 means infinite)
2526 # This is useful when the aperturedb server is not ready yet.
2627 retry_max_attempts : int = 3
28+ token : str = None
29+ user_keys : dict = None
2730
2831 def __repr__ (self ) -> str :
2932 mode = "REST" if self .use_rest else "TCP"
30- return f"[{ self .host } { self .port } { self .username } { mode } { self .use_ssl }
33+ auth_mode = "token" if self .token is not None else "password"
34+ return f"[{ self .host } { self .port } { self .username } { mode } { self .use_ssl } { auth_mode }
3135
3236 def deflate (self ) -> list :
3337 deflate_version = 1
@@ -41,22 +45,96 @@ def deflate(self) -> list:
4145 encoded_key = b64encode (simplified .encode ('utf-8' )).decode ('utf-8' )
4246 return encoded_key
4347
48+ def has_user_keys (self ) -> bool :
49+ return self .user_keys is not None
50+
51+ def add_user_key (self , user , key ):
52+ if self .user_keys is None :
53+ self .user_keys = dict ()
54+ if not user in self .user_keys :
55+ self .user_keys [user ] = []
56+ self .user_keys [user ].insert (0 , key )
57+
58+ def get_user_key (self , for_user : str ) -> str :
59+ if self .user_keys is None or not (for_user in self .user_keys ) \
60+ or len (self .user_keys [for_user ]) == 0 :
61+ return None
62+ return self .user_keys [for_user ][0 ]
63+
64+ def set_user_keys (self , keys : dict ) -> None :
65+ self .user_keys = keys
66+
67+ @classmethod
68+ def create_web_token (cls , host : str , port : int , token_string : str ) -> None :
69+ if token_string .startswith ("adbp_" ):
70+ token_string = token_string [5 :]
71+
72+ if host .endswith (APERTURE_CLOUD ):
73+ host = host [:- 1 * len (APERTURE_CLOUD )]
74+ m = re .match ("(.*)\.farm(\d+)$" , host )
75+ if m is not None :
76+ host = "{}.{}" .format (m .group (1 ), int (m .group (2 )))
77+
78+ host = "a://{}" .format (host )
79+
80+ if port == 55555 :
81+ web_token_json = [2 , host , token_string ]
82+ else :
83+ web_token_json = [2 , host , port , token_string ]
84+ simplified = json .dumps (web_token_json )
85+ encoded = b64encode (simplified .encode ('utf-8' )).decode ('utf-8' )
86+ return encoded
87+
4488 @classmethod
45- def reinflate (cls , encoded_key : list ) -> object :
46- decoded_key = b64decode (encoded_key .encode ('utf-8' ))
47- as_list = json .loads (decoded_key .decode ('utf-8' ))
48- if as_list [0 ] != 1 :
49- raise ValueError (f"version identifier of configuration was"
50- "{as_list[0]}, which is not supported" )
51- host , port , username , password , name , use_ssl , \
52- use_rest , use_keepalive , retry_interval_seconds , \
53- retry_max_attempts = as_list [1 :]
54- if host .startswith (AD_CLOUD_SCHEME ):
55- host = host [len (AD_CLOUD_SCHEME ):] + APERTURE_CLOUD
56- use_ssl = bool (use_ssl )
57- use_rest = bool (use_rest )
58- use_keepaliave = bool (use_keepalive )
59- c = Configuration (
60- host , port , username , password , name , use_ssl , use_rest , use_keepalive ,
61- retry_interval_seconds , retry_max_attempts )
89+ def reinflate (cls , encoded_str : list ) -> object :
90+ try :
91+ decoded = b64decode (encoded_str .encode ('utf-8' ))
92+ as_list = json .loads (decoded .decode ('utf-8' ))
93+ except :
94+ raise Exception (
95+ "Unable to make configuration from the provided string" )
96+ version = as_list [0 ]
97+ if version not in 1 , 2 ):
98+ raise ValueError ("version identifier of configuration was"
99+ f"g{ as_list [0 ]} )
100+ if version == 1 :
101+ host , port , username , password , name , use_ssl , \
102+ use_rest , use_keepalive , retry_interval_seconds , \
103+ retry_max_attempts = as_list [1 :]
104+ if host .startswith (AD_CLOUD_SCHEME ):
105+ host = host [len (AD_CLOUD_SCHEME ):] + APERTURE_CLOUD
106+ use_ssl = bool (use_ssl )
107+ use_rest = bool (use_rest )
108+ use_keepaliave = bool (use_keepalive )
109+ c = Configuration (
110+ host , port , username , password , name , use_ssl , use_rest , use_keepalive ,
111+ retry_interval_seconds , retry_max_attempts )
112+ elif version == 2 :
113+ host = as_list [1 ]
114+ if host .startswith ("a://" ):
115+ m = re .match ("a://([^.]*)\.(\d+)" , host )
116+ host = "{}.farm{:04d}.cloud.aperturedata.io" .format (
117+ m .group (1 ), int (m .group (2 )))
118+
119+ cur_arg = 2
120+ # second arg is port
121+ if isinstance (as_list [2 ], int ):
122+ cur_arg = 3
123+ else :
124+ port = 55555
125+
126+ name = "default"
127+
128+ username = None
129+ password = None
130+ token = None
131+ # if 1 argument left, treat as token.
132+ if len (as_list ) - cur_arg == 1 :
133+ token = "adbp_" + as_list [cur_arg ]
134+ else :
135+ username = as_list [cur_arg ]
136+ password = as_list [cur_arg + 1 ]
137+
138+ c = Configuration (host , port , username ,
139+ password , name , token = token )
62140 return c
0 commit comments