update ver_status on each new fetch data

Signed-off-by: Vivek Kumar Sahu <vivekkumarsahu650@gmail.com>

Author: viveksahu26

lynkctx.py

@@ -4,9 +4,7 @@
 import base64
 import requests
 
-# INTERLYNK_API_URL = 'https://api.interlynk.io/lynkapi'
-
-INTERLYNK_API_URL = 'http://localhost:3000/lynkapi'
+INTERLYNK_API_URL = 'https://api.interlynk.io/lynkapi'
 
 INTERLYNK_API_TIMEOUT = 100
 
@@ -229,20 +227,16 @@ def resolve_env(self):
 
     def resolve_ver(self):
         env = self.env or 'default'
-        self.data = self._fetch_context()
-
         if not self.ver_id:
             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:
                             for ver in env['versions']:
-                                if ver.get('primaryComponent') and ver['primaryComponent'].get('version') == self.ver:
+                                if ver['primaryComponent']['version'] == self.ver:
                                     self.ver_id = ver['id']
                                     self.ver_status = self.vuln_status_to_status(
-                                        ver['vulnRunStatus']
-                                    )
-
+                                        ver['vulnRunStatus'])
         empty_ver = False
         if not self.ver:
             for product in self.data.get('data', {}).get('organization', {}).get('productNodes', {}).get('products', []):
@@ -251,13 +245,11 @@ def resolve_ver(self):
                         if env['id'] == self.env_id:
                             for ver in env['versions']:
                                 if ver['id'] == self.ver_id:
-                                    if ver.get('primaryComponent'):
-                                        self.ver = ver['primaryComponent'].get('version')
+                                    self.ver = ver['primaryComponent']['version']
                                     if not self.ver:
                                         empty_ver = True
                                     self.ver_status = self.vuln_status_to_status(
-                                        ver['vulnRunStatus']
-                                    )
+                                        ver['vulnRunStatus'])
 
         return (empty_ver or self.ver) and self.ver_id
 
@@ -296,6 +288,11 @@ def versions(self):
     def status(self):
         self.resolve_ver()
         return self.ver_status
+    
+    def live_status(self):
+        self.resolve_ver_status()
+        return self.ver_status
+
 
     def download(self):
         logging.debug("Downloading SBOM for environment ID %s, sbom ID %s",
@@ -447,3 +444,60 @@ def vuln_status_to_status(self, status):
             result_dict['labelingStatus'] = 'COMPLETED'
             result_dict['automationStatus'] = 'COMPLETED'
         return result_dict
+
+
+    def resolve_ver_status(self):
+        """
+        Resolve version ID (ver_id) and version (ver) for the current context.
+        """
+        self.data = self._fetch_context()
+
+        # ver_id is present
+        if self.ver_id:
+            self._update_ver_status()
+            return self.ver_id
+
+        # ver_id is not present
+        self._resolve_ver_id()
+
+        # ver is not present
+        if not self.ver:
+            self._resolve_ver_value()
+
+        return self.ver_id and self.ver
+
+    def _update_ver_status(self):
+        """Update the status of the version based on ver_id."""
+        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:
+                        for ver in env['versions']:
+                            if ver['id'] == self.ver_id:
+                                self.ver_status = self.vuln_status_to_status(ver['vulnRunStatus'])
+
+    def _resolve_ver_id(self):
+        """Resolve ver_id based on ver."""
+        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:
+                        for ver in env['versions']:
+                            if ver.get('primaryComponent') and ver['primaryComponent'].get('version') == self.ver:
+                                self.ver_id = ver['id']
+                                self.ver_status = self.vuln_status_to_status(ver['vulnRunStatus'])
+
+    def _resolve_ver_value(self):
+        """Resolve the version value (ver) based on ver_id."""
+        empty_ver = False
+        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:
+                        for ver in env['versions']:
+                            if ver['id'] == self.ver_id:
+                                if ver.get('primaryComponent'):
+                                    self.ver = ver['primaryComponent'].get('version')
+                                if not self.ver:
+                                    empty_ver = True
+        return empty_ver

pylynk.py

@@ -213,7 +213,7 @@ def download_sbom(lynk_ctx):
     return 0
 
 
-def upload_sbom(lynk_ctx, sbom_file):
+def upload_sbom(lynk_ctx, sbom_file, download):
     """
     Upload SBOM to the lynk_ctx.
 
@@ -224,22 +224,27 @@ def upload_sbom(lynk_ctx, sbom_file):
     Returns:
         The result of the upload operation.
     """
-    # return lynk_ctx.upload(sbom_file)
     upload_result = lynk_ctx.upload(sbom_file)
     if upload_result != 0:
         return 1
 
-    if True:
-        while True:
-            status = lynk_ctx.status()
+    if download:
+        max_retries = 5  # Set the maximum number of retries
+        retries = 0
+        while retries < max_retries:
+            status = lynk_ctx.live_status()
 
             if status.get('automationStatus') == "COMPLETED":
                 download_sbom(lynk_ctx)
                 break
             else:
-                print("Waiting for automation status to complete...")
                 time.sleep(5)
+                retries += 1
 
+        if retries == max_retries:                                                                                                                                                                 
+            print("Error: automationStatus could not be completed within the maximum retry limit.")
+            return 1
+        
     return 0
 
 
@@ -306,6 +311,7 @@ def setup_args():
     upload_parser.add_argument("--token",
                                required=False,
                                help="Security token")
+    upload_parser.add_argument("--download", action="store_true", help="Download SBOM after upload (default: False)")
 
     download_parser = subparsers.add_parser("download", help="Download SBOM")
     download_group = download_parser.add_mutually_exclusive_group(
@@ -390,7 +396,8 @@ def main() -> int:
     elif args.subcommand == "status":
         print_status(lynk_ctx, fmt_json)
     elif args.subcommand == "upload":
-        upload_sbom(lynk_ctx, args.sbom)
+        download_flag = getattr(args, 'download', False)
+        upload_sbom(lynk_ctx, args.sbom, download_flag)
     elif args.subcommand == "download":
         download_sbom(lynk_ctx)
     else: