Merge pull request #809 from Ostorlab/feature/add-option-to-order-vul…

…nz-by-risk_id_title

Add option to order vulnerabilities by id, title or risk_rating

src/ostorlab/cli/vulnz/list/list.py

@@ -22,12 +22,21 @@
     help="Filter vulnerabilities by risk ratings. Accept comma-separated ratings.",
 )
 @click.option("--search", "-sh", help="Search in all content of the vulnerabilities.")
+@click.option(
+    "--order-by",
+    "-o",
+    "order_by",
+    type=click.Choice(["risk_rating", "title", "id"], case_sensitive=False),
+    default="risk_rating",
+    help="Order vulnerabilities by specified field. Defaults to risk_rating.",
+)
 @click.pass_context
 def list_cli(
     ctx: click.core.Context,
     scan_id: int,
     risk_rating: Optional[str],
     search: Optional[str],
+    order_by: str,
 ) -> None:
     """CLI command to list vulnerabilities for a scan."""
     runtime_instance = ctx.obj["runtime"]
@@ -36,4 +45,5 @@ def list_cli(
         scan_id=scan_id,
         filter_risk_rating=risk_rating.strip().split(",") if risk_rating else None,
         search=search,
+        order_by=order_by,
     )

src/ostorlab/runtimes/cloud/runtime.py

@@ -252,6 +252,7 @@ def list_vulnz(
         number_elements: int = 10,
         filter_risk_rating: Optional[List[str]] = None,
         search: Optional[str] = None,
+        order_by: Optional[str] = None,
     ) -> None:
         """List vulnz from the cloud using and render them in a table.
 

src/ostorlab/runtimes/local/runtime.py

@@ -635,8 +635,9 @@ def install(self, docker_client: Optional[docker.DockerClient] = None) -> None:
     def list_vulnz(
         self,
         scan_id: int,
-        filter_risk_rating: Optional[List[str]],
-        search: Optional[str],
+        filter_risk_rating: Optional[List[str]] = None,
+        search: Optional[str] = None,
+        order_by: str = "risk_rating",
     ) -> None:
         try:
             with models.Database() as session:
@@ -658,7 +659,22 @@ def list_vulnz(
                         )
                     )
 
-                vulnerabilities = query.order_by(models.Vulnerability.title).all()
+                if order_by == "risk_rating":
+                    case_ordering = case(
+                        [
+                            (models.Vulnerability.risk_rating == rating, order)
+                            for rating, order in risk_rating.RATINGS_ORDER.items()
+                        ],
+                        else_=len(risk_rating.RATINGS_ORDER),
+                    )
+                    vulnerabilities = query.order_by(
+                        case_ordering, models.Vulnerability.title
+                    ).all()
+                elif order_by == "title":
+                    vulnerabilities = query.order_by(models.Vulnerability.title).all()
+                elif order_by == "id":
+                    vulnerabilities = query.order_by(models.Vulnerability.id).all()
+
             vulnz_list = []
             for vulnerability in vulnerabilities:
                 vulnerability_location = vulnerability.location or ""

src/ostorlab/serve_app/types.py

@@ -712,15 +712,17 @@ def resolve_progress(self: models.Scan, info: graphql_base.ResolveInfo) -> str:
 
         return self.progress.name
 
-    def resolve_risk_rating(self: models.Scan, info: graphql_base.ResolveInfo) -> str:
+    def resolve_risk_rating(
+        self: models.Scan, info: graphql_base.ResolveInfo
+    ) -> Optional[str]:
         """Resolve risk rating query.
         Args:
             self (models.Scan): The scan object.
             info (graphql_base.ResolveInfo): GraphQL resolve info.
         Returns:
             str: The risk rating of the scan.
         """
-        return self.risk_rating.name
+        return self.risk_rating.name if self.risk_rating is not None else None
 
     def resolve_assets(
         self,

tests/cli/vulnz/list/test_list_vulnz.py

@@ -475,3 +475,112 @@ def testOstorlabVulnzListCLI_whenFilterBySearchAndRuntimeIsCloud_showsCorrectRes
         all(word in result_tech_detail.output for word in result_tech_detail_keywords)
         is True
     )
+
+
+def testOstorlabVulnzListCLI_whenListVulnz_showsVulnzOrderedByRiskRatingByDefault(
+    scan_multiple_vulnz_different_risk_ratings: models.Scan,
+    mocker: plugin.MockerFixture,
+) -> None:
+    """Test oxo vulnz list command orders vulnerabilities by risk rating."""
+    runner = CliRunner()
+    table_mock = mocker.patch("ostorlab.cli.console.Console.table")
+
+    result = runner.invoke(
+        rootcli.rootcli,
+        ["vulnz", "list", "-s", str(scan_multiple_vulnz_different_risk_ratings.id)],
+    )
+
+    assert result.exception is None
+    risk_ratings = [
+        vuln.get("risk_rating") for vuln in table_mock.call_args_list[0][1].get("data")
+    ]
+    assert risk_ratings == [
+        "[bold bright_white on #F55246]High[/]",
+        "[bold bright_white on #FF9800]Medium[/]",
+        "[bold bright_white on #FF9800]Medium[/]",
+        "[bold bright_white on #FDDB45]Low[/]",
+    ]
+
+
+def testOstorlabVulnzListCLI_whenListVulnzOrderByID_showsVulnzOrderedByID(
+    scan_multiple_vulnz_different_risk_ratings: models.Scan,
+    mocker: plugin.MockerFixture,
+) -> None:
+    """Test oxo vulnz list command orders vulnerabilities by ID."""
+    runner = CliRunner()
+    table_mock = mocker.patch("ostorlab.cli.console.Console.table")
+
+    result = runner.invoke(
+        rootcli.rootcli,
+        [
+            "vulnz",
+            "list",
+            "-s",
+            str(scan_multiple_vulnz_different_risk_ratings.id),
+            "-o",
+            "id",
+        ],
+    )
+
+    assert result.exception is None
+    ids = [vuln.get("id") for vuln in table_mock.call_args_list[0][1].get("data")]
+    assert ids == ["1", "2", "3", "4"]
+
+
+def testOstorlabVulnzListCLI_whenListVulnzOrderByTitle_showsVulnzOrderedByTitle(
+    scan_multiple_vulnz_different_risk_ratings: models.Scan,
+    mocker: plugin.MockerFixture,
+) -> None:
+    """Test oxo vulnz list command orders vulnerabilities by Title."""
+    runner = CliRunner()
+    table_mock = mocker.patch("ostorlab.cli.console.Console.table")
+
+    result = runner.invoke(
+        rootcli.rootcli,
+        [
+            "vulnz",
+            "list",
+            "-s",
+            str(scan_multiple_vulnz_different_risk_ratings.id),
+            "-o",
+            "title",
+        ],
+    )
+
+    assert result.exception is None
+    titles = [vuln.get("title") for vuln in table_mock.call_args_list[0][1].get("data")]
+    assert titles == [
+        "vulnerability 1",
+        "vulnerability 2",
+        "vulnerability 3",
+        "vulnerability 4",
+    ]
+
+
+def testOstorlabVulnzListCLI_whenListVulnzOrderByInvalidOption_showsErrorMessage(
+    scan_multiple_vulnz_different_risk_ratings: models.Scan,
+) -> None:
+    """Test oxo vulnz list command orders vulnerabilities by Title."""
+    runner = CliRunner()
+
+    result = runner.invoke(
+        rootcli.rootcli,
+        [
+            "vulnz",
+            "list",
+            "-s",
+            str(scan_multiple_vulnz_different_risk_ratings.id),
+            "-o",
+            "invalid",
+        ],
+    )
+
+    assert result.exception is not None
+    assert (
+        result.output.replace("\r\n", "\n")
+        == """Usage: rootcli vulnz list [OPTIONS]
+Try 'rootcli vulnz list --help' for help.
+
+Error: Invalid value for '--order-by' / '-o': 'invalid' is not one of 'risk_rating', 'title', 'id'.
+"""
+    )

tests/conftest.py

@@ -1509,3 +1509,64 @@ def multiple_assets_scan_bytes() -> bytes:
     """Returns a dummy zip file."""
     zip_path = pathlib.Path(__file__).parent / "files" / "multiple_assets_scan.zip"
     return zip_path.read_bytes()
+
+
+@pytest.fixture
+def scan_multiple_vulnz_different_risk_ratings(
+    clean_db: None, mocker: plugin.MockerFixture, db_engine_path: str
+) -> models.Scan:
+    mocker.patch.object(models, "ENGINE_URL", db_engine_path)
+    create_scan_db = models.Scan.create("test")
+    models.Vulnerability.create(
+        title="vulnerability 1",
+        short_description="vulnerability 1",
+        description="vulnerability 1",
+        recommendation="Consider fixing soon",
+        technical_detail="example=$input",
+        risk_rating="MEDIUM",
+        cvss_v3_vector="5:6:7",
+        dna="12347",
+        location={},
+        scan_id=create_scan_db.id,
+        references=[],
+    )
+    models.Vulnerability.create(
+        title="vulnerability 2",
+        short_description="vulnerability 2",
+        description="vulnerability 2",
+        recommendation="Fix immediately",
+        technical_detail="example=$input",
+        risk_rating="HIGH",
+        cvss_v3_vector="5:6:7",
+        dna="12345",
+        location={},
+        scan_id=create_scan_db.id,
+        references=[],
+    )
+    models.Vulnerability.create(
+        title="vulnerability 3",
+        short_description="vulnerability 3",
+        description="vulnerability 3",
+        recommendation="Monitor the situation",
+        technical_detail="example=$input",
+        risk_rating="LOW",
+        cvss_v3_vector="5:6:7",
+        dna="12346",
+        location={},
+        scan_id=create_scan_db.id,
+        references=[],
+    )
+    models.Vulnerability.create(
+        title="vulnerability 4",
+        short_description="vulnerability 4",
+        description="vulnerability 4",
+        recommendation="Consider fixing soon",
+        technical_detail="example=$input",
+        risk_rating="MEDIUM",
+        cvss_v3_vector="5:6:7",
+        dna="12347",
+        location={},
+        scan_id=create_scan_db.id,
+        references=[],
+    )
+    return create_scan_db