support download option for upload command

Signed-off-by: Vivek Kumar Sahu <[email protected]>

Author: viveksahu26

lynkctx.py

@@ -226,6 +226,7 @@ def resolve_env(self):
 
     def resolve_ver(self):
         env = self.env or 'default'
+        print("self.ver_id: ", self.ver_id)
         if not self.ver_id:
             for product in self.data.get('data', {}).get('organization', {}).get('productNodes', {}).get('products', []):
                 if product['id'] == self.prod_id:
@@ -252,6 +253,46 @@ def resolve_ver(self):
 
         return (empty_ver or self.ver) and self.ver_id
 
+    def resolve_latest_ver(self):
+        """
+        Resolve and set the version ID for the latest uploaded SBOM version
+        based on the most recent updatedAt timestamp.
+        """
+        # fetch the context data
+        self.data = self._fetch_context()
+        if not self.data:
+            print("Failed to refresh context data")
+            return False
+
+        print("Attempting to resolve the latest version information for Product ID:", self.prod_id, "Environment ID:", self.env_id)
+
+        # retry to handle propagation delays
+        for attempt in range(5):
+            print(f"Attempt {attempt + 1}: Resolving the latest version ID based on updatedAt")
+
+            for product in self.data.get('data', {}).get('organization', {}).get('productNodes', {}).get('products', []):
+                if product['id'] == self.prod_id:
+                    for env in product['environments']:
+                        if env['id'] == self.env_id:
+
+                            # sort the versions by updatedAt in descending order
+                            versions_sorted = sorted(env['versions'], key=lambda v: v['updatedAt'], reverse=True)
+                            latest_version = versions_sorted[0] if versions_sorted else None
+                            
+                            if latest_version:
+                                self.ver_id = latest_version['id']
+                                self.ver = latest_version['primaryComponent']['version']
+                                self.ver_status = self.vuln_status_to_status(latest_version['vulnRunStatus'])
+                                print(f"Resolved version ID: {self.ver_id}")
+                                return True
+
+            print("Latest version ID not found. Retrying...")
+            time.sleep(2)
+
+        print("Error: Unable to resolve the latest SBOM version information after multiple attempts.")
+        return False
+
+
     def __repr__(self):
         from pprint import pformat
         return pformat(vars(self), indent=1, width=1)
@@ -286,6 +327,7 @@ def versions(self):
 
     def status(self):
         self.resolve_ver()
+        print("self.ver_status: ", self.ver_status)
         return self.ver_status
 
     def download(self):
@@ -334,17 +376,17 @@ def download(self):
                           response.status_code)
 
     def upload(self, sbom_file):
-        if os.path.isfile(sbom_file) is False:
+        if not os.path.isfile(sbom_file):
             print(f"SBOM File not found: {sbom_file}")
-            return
+            return None
 
         if not self.prod_id:
             print(f"Product not found: {self.prod}")
-            return
+            return None
 
         if not self.env_id:
             print(f"Environment not found: {self.env}")
-            return
+            return None
 
         logging.debug("Uploading SBOM to product ID %s", self.prod_id)
 
@@ -367,26 +409,66 @@ def upload(self, sbom_file):
             with open(sbom_file, 'rb') as sbom:
                 files_map = {'0': sbom}
                 response = requests.post(self.api_url,
-                                         headers=headers,
-                                         data=form_data,
-                                         files=files_map,
-                                         timeout=INTERLYNK_API_TIMEOUT)
+                                        headers=headers,
+                                        data=form_data,
+                                        files=files_map,
+                                        timeout=INTERLYNK_API_TIMEOUT)
                 if response.status_code == 200:
                     resp_json = response.json()
-                    errors = resp_json.get('data', {}).get(
-                        'sbomUpload', {}).get('errors')
+                    errors = resp_json.get('data', {}).get('sbomUpload', {}).get('errors')
                     if errors:
                         print(f"Error uploading sbom: {errors}")
-                        return 1
+                        return None
+
                     print('Uploaded successfully')
                     logging.debug("SBOM Uploading response: %s", response.text)
-                    return 0
+                    return True  # Return True to indicate a successful upload
+
                 logging.error("Error uploading sbom: %d", response.status_code)
         except requests.exceptions.RequestException as ex:
             logging.error("RequestException: %s", ex)
         except FileNotFoundError as ex:
             logging.error("FileNotFoundError: %s", ex)
-        return 1
+
+        return None
+
+
+    def _fetch_latest_version_id(self):
+        """
+        Fetch the latest version ID for the current product and environment.
+        """
+        headers = {
+            "Authorization": "Bearer " + self.token
+        }
+
+        request_data = {
+            "query": QUERY_PRODUCTS_LIST,
+            "variables": {"first": 1},  # Fetch the latest
+        }
+
+        try:
+            response = requests.post(self.api_url,
+                                    headers=headers,
+                                    json=request_data,
+                                    timeout=INTERLYNK_API_TIMEOUT)
+            if response.status_code == 200:
+                response_data = response.json()
+                products = response_data.get('data', {}).get('organization', {}).get('productNodes', {}).get('products', [])
+                for product in products:
+                    if product['id'] == self.prod_id:
+                        for env in product['environments']:
+                            if env['id'] == self.env_id:
+                                # Assuming the latest version is at the end of the list
+                                latest_version = env['versions'][-1]
+                                return latest_version['id'] if latest_version else None
+
+                logging.error("Error: Version ID not found for the specified product and environment.")
+            else:
+                logging.error("Failed to retrieve latest version ID. Status code: %s", response.status_code)
+        except requests.exceptions.RequestException as ex:
+            logging.error("RequestException: %s", ex)
+        return None
+
 
     def vuln_status_to_status(self, status):
         result_dict = dict()
@@ -396,18 +478,21 @@ def vuln_status_to_status(self, status):
         result_dict['automationStatus'] = 'UNKNOWN'
         result_dict['vulnScanStatus'] = 'UNKNOWN'
         if status == 'NOT_STARTED':
+            print("NOT_STARTED")
             result_dict['vulnScanStatus'] = 'NOT_STARTED'
             result_dict['checksStatus'] = 'NOT_STARTED'
             result_dict['policyStatus'] = 'NOT_STARTED'
             result_dict['labelingStatus'] = 'NOT_STARTED'
             result_dict['automationStatus'] = 'NOT_STARTED'
         elif status == 'IN_PROGRESS':
+            print("IN_PROGRESS")
             result_dict['vulnScanStatus'] = 'IN_PROGRESS'
             result_dict['checksStatus'] = 'COMPLETED'
             result_dict['policyStatus'] = 'COMPLETED'
             result_dict['labelingStatus'] = 'COMPLETED'
             result_dict['automationStatus'] = 'COMPLETED'
         elif status == 'FINISHED':
+            print("FINISHED")
             result_dict['vulnScanStatus'] = 'COMPLETED'
             result_dict['checksStatus'] = 'COMPLETED'
             result_dict['policyStatus'] = 'COMPLETED'

pylynk.py

@@ -212,19 +212,48 @@ def download_sbom(lynk_ctx):
 
     return 0
 
-
-def upload_sbom(lynk_ctx, sbom_file):
+def upload_sbom(lynk_ctx, sbom_file, download=False):
     """
-    Upload SBOM to the lynk_ctx.
+    Upload SBOM to the lynk_ctx and optionally download it if specified.
 
     Args:
         lynk_ctx: The lynk context object.
         sbom_file: The path to the SBOM file.
+        download: Boolean to indicate if download is requested after upload.
 
     Returns:
         The result of the upload operation.
     """
-    return lynk_ctx.upload(sbom_file)
+    upload_result = lynk_ctx.upload(sbom_file)
+    if not upload_result:
+        print("SBOM upload failed.")
+        return 1
+
+    print("SBOM uploaded successfully.")
+    print(f"Product ID: {lynk_ctx.prod_id}, Environment ID: {lynk_ctx.env_id}")
+
+    # get the version ID after upload
+    if not lynk_ctx.resolve_latest_ver():
+        print("Error: Unable to resolve SBOM version information.")
+        return 1
+
+    if download:
+        while True:
+            print("Checking SBOM status for automation completion...")
+
+            # retrieve status for specific version ID
+            status = lynk_ctx.status()
+            print("Current Status:", status)
+
+            if status.get('automationStatus') == "COMPLETED":
+                print("Automation status completed. Initiating SBOM download.")
+                download_sbom(lynk_ctx)
+                break
+            else:
+                print("Waiting for automation status to complete...")
+                time.sleep(5)
+
+    return 0
 
 
 def add_output_format_group(parser):
@@ -290,6 +319,8 @@ def setup_args():
     upload_parser.add_argument("--token",
                                required=False,
                                help="Security token")
+    upload_parser.add_argument("--download", required=False, help="Download SBOM after upload")
+
 
     download_parser = subparsers.add_parser("download", help="Download SBOM")
     download_group = download_parser.add_mutually_exclusive_group(
@@ -374,7 +405,7 @@ def main() -> int:
     elif args.subcommand == "status":
         print_status(lynk_ctx, fmt_json)
     elif args.subcommand == "upload":
-        upload_sbom(lynk_ctx, args.sbom)
+        upload_sbom(lynk_ctx, args.sbom, download=args.download)
     elif args.subcommand == "download":
         download_sbom(lynk_ctx)
     else: