PullRequest: 46 fix: set a longer validity period for the authentication information of the method upload_pkg.

Merge branch fix_upload_pkd of git@code.alipay.com:oceanbase/OBShell-SDK-Python.git into master
https://code.alipay.com/oceanbase/OBShell-SDK-Python/pull_requests/46?tab=comment

Signed-off-by: 玉楼 <rongfeng.frf@alibaba-inc.com>


* fix: set a longer validity period for the authentication information of

Author: Junkrat77

obshell/auth/password.py

@@ -34,7 +34,7 @@
 class PasswordAuth(base.Auth):
     """Password-based authentication method."""
 
-    def __init__(self, password: str = "", version=None) -> None:
+    def __init__(self, password: str = "", version=None, lifetime=60) -> None:
         """Initialize a new PasswordAuth instance.
 
         Args:
@@ -49,10 +49,14 @@ def __init__(self, password: str = "", version=None) -> None:
 
                 - "v1": supported by OBShell version 4.2.2.0.
                 - "v2": supported by OBShell version 4.2.3.0 or later.
+            lifetime (int, optional): 
+                lifetime of the authentication information in reqeust header.
+                Defalut to 60 second.
         """
         super().__init__(base.AuthType.PASSWORD,
                          [base.AuthVersion.V1, base.AuthVersion.V2])
         self.password = password
+        self.lifetime = lifetime
         if version is not None:
             if version not in _AUTHS_VERSION:
                 raise ValueError("Version not supported")
@@ -63,15 +67,16 @@ def auth(self, request) -> None:
             version = self.get_version()
             if version not in _AUTHS:
                 raise base.AuthError(f"Unsupported auth version: {version}")
-            self._method = _AUTHS[version](self.password)
+            self._method = _AUTHS[version](self.password, self.lifetime)
         self._method.auth(request)
 
 
 class PasswordAuthMethod:
 
-    def __init__(self, password: str) -> None:
+    def __init__(self, password: str, lifetime: int) -> None:
         self.password = password
         self.pk = None
+        self.lifetime = lifetime
         self.check_identity = False
 
     def reset(self) -> None:
@@ -99,7 +104,7 @@ def auth(self, req: requests.Request) -> None:
         self._check(req.server)
         self._init_pk(req.server)
         auth_json = json.dumps(
-            {'password': self.password, 'ts': int(time.time()) + 5})
+            {'password': self.password, 'ts': int(time.time() + self.lifetime)})
         key = RSA.import_key(base64.b64decode(self.pk))
         cipher = PKCS1_cipher.new(key)
         req.headers['X-OCS-Auth'] = base64.b64encode(
@@ -132,22 +137,13 @@ def encrypt_header(self, headers: str) -> str:
     def auth(self, req: requests.Request) -> None:
         self._check(req.server)
         self._init_pk(req.server)
+
+        # encrypt body before build header.
         aes_key = get_random_bytes(16)
         aes_iv = get_random_bytes(16)
-        uri = urlparse(req.url).path if not urlparse(
-            req.url).query else urlparse(req.url).path + "?" + urlparse(req.url).query
-        headers = {
-            'auth': self.password,
-            'ts': str(int(time.time()) + 5),
-            'uri': uri,
-            'keys': base64.b64encode(aes_key+aes_iv).decode('utf-8')
-        }
-        req.headers['X-OCS-Header'] = self.encrypt_header(headers)
-
         cipher = AES.new(aes_key, AES.MODE_CBC, aes_iv)
         if not req.original_data:
             req.original_data = req.data
-
         if req.original_data:
             body = None
             if isinstance(req.original_data, dict):
@@ -162,6 +158,16 @@ def auth(self, req: requests.Request) -> None:
             req.data = base64.b64encode(
                 cipher.encrypt(pad(bytes(body), AES.block_size))
             ).decode('utf8')
+
+        uri = urlparse(req.url).path if not urlparse(
+            req.url).query else urlparse(req.url).path + "?" + urlparse(req.url).query
+        headers = {
+            'auth': self.password,
+            'ts': str(int(time.time()) + self.lifetime),
+            'uri': uri,
+            'keys': base64.b64encode(aes_key+aes_iv).decode('utf-8')
+        }
+        req.headers['X-OCS-Header'] = self.encrypt_header(headers)
         return
 
 

obshell/service/client_v1.py

@@ -664,11 +664,19 @@ def upload_pkg(self, pkg_path: str) -> UpgradePkgInfo:
         Raises:
             OBShellHandleError: error message return by OBShell server.
         """
-        req = self.create_request("/api/v1/upgrade/package", "POST")
         data, headers = self._parse_pkg(pkg_path)
-        req.data = data
-        req.headers = headers
-        return self._handle_ret_request(req, UpgradePkgInfo)
+        auth = self._get_auth()
+        try:
+            if auth.type == AuthType.PASSWORD:
+                copied_auth = copy.deepcopy(auth)
+                copied_auth.lifetime = 600
+                self._set_auth(copied_auth)
+            req = self.create_request("/api/v1/upgrade/package", "POST")
+            req.data = data
+            req.headers = headers
+            self._handle_ret_request(req, UpgradePkgInfo)
+        finally:
+            self._set_auth(auth)
 
     def upgrade_agent_check(
             self, version: str, release: str, upgrade_dir=None) -> task.DagDetailDTO:
@@ -2309,21 +2317,21 @@ def get_tenant_restore_overview(self, tenant_name: str) -> ob.RestoreOverview:
         return self._handle_ret_request(req, ob.RestoreOverview)
 
     def get_restore_windows(
-        self, 
-        data_backup_uri: str, 
+        self,
+        data_backup_uri: str,
         archive_log_uri: str = None,
-        ):
+    ):
         """Get windows during which the tenant can be restored.
-        
+
         Args:
             data_backup_uri (str): Complete destination path for data backups.
             archive_log_uri (str, optional): Destination path for log archives.
         """
-        
+
         data = {
             'data_backup_uri': data_backup_uri,
         }
         if archive_log_uri is not None:
             data['archive_log_uri'] = archive_log_uri
         req = self.create_request("/api/v1/restore/windows", "GET", data)
-        return self._handle_ret_request(req, ob.RestoreWindows)
+        return self._handle_ret_request(req, ob.RestoreWindows)