Merge pull request #692 from Ostorlab/feature/ensure_oxo_schemas_are_…

…similar

Ensure Oxo serve & RE_Oxo schemas are similar.

.github/workflows/pytest-linux.yml

@@ -41,6 +41,8 @@ jobs:
         pip install .[serve]
         python -m pip install -e .
     - name: Running tests with pytest.
+      env:
+        RE_OXO_API_KEY: ${{ secrets.RE_OXO_API_KEY }}
       run: |
         set -o pipefail
         pytest -m "not docker and not nats"

.github/workflows/pytest-windows.yml

@@ -41,5 +41,7 @@ jobs:
         pip install .[serve]
         python -m pip install -e .
     - name: Running tests with pytest.
+      env:
+        RE_OXO_API_KEY: ${{ secrets.RE_OXO_API_KEY }}
       run: |
         pytest -m "not docker and not nats"

src/ostorlab/serve_app/types.py

@@ -12,6 +12,7 @@
 
 from ostorlab.runtimes.local.models import models
 from ostorlab.serve_app import common
+from ostorlab.utils import risk_rating as utils_rik_rating
 
 DEFAULT_NUMBER_ELEMENTS = 15
 RISK_RATINGS_ORDER = {
@@ -59,11 +60,18 @@ class OxoKnowledgeBaseVulnerabilityType(graphene.ObjectType):
     recommendation = graphene.String()
 
 
+OxoRiskRatingEnum = graphene.Enum(
+    "OxoRiskRating",
+    [(risk.name.upper(), i) for i, risk in enumerate(utils_rik_rating.RiskRating)],
+)
+
+
 class OxoVulnerabilityType(graphene_sqlalchemy.SQLAlchemyObjectType):
     """SQLAlchemy object type for a vulnerability."""
 
     detail = graphene.Field(OxoKnowledgeBaseVulnerabilityType, required=False)
     cvss_v3_base_score = graphene.Float(required=False)
+    risk_rating = graphene.Field(OxoRiskRatingEnum, required=False)
 
     class Meta:
         """Meta class for the vulnerability object type."""
@@ -120,6 +128,15 @@ def resolve_detail(
                 ],
             )
 
+    def resolve_risk_rating(
+        self: models.Vulnerability, info: graphql_base.ResolveInfo
+    ) -> Optional[OxoRiskRatingEnum]:
+        """Resolve risk rating of vulnerability"""
+        try:
+            return OxoRiskRatingEnum[self.risk_rating.name]
+        except KeyError:
+            return None
+
 
 class OxoVulnerabilitiesType(graphene.ObjectType):
     """Graphene object type for a list of vulnerabilities."""
@@ -131,7 +148,7 @@ class OxoVulnerabilitiesType(graphene.ObjectType):
 class OxoAggregatedKnowledgeBaseVulnerabilityType(graphene.ObjectType):
     """Graphene object type for an aggregated knowledge base vulnerability."""
 
-    highest_risk_rating = graphene.Field(common.RiskRatingEnum)
+    highest_risk_rating = graphene.Field(OxoRiskRatingEnum)
     highest_cvss_v3_vector = graphene.String()
     highest_cvss_v3_base_score = graphene.Float()
     kb = graphene.Field(OxoKnowledgeBaseVulnerabilityType)
@@ -143,6 +160,12 @@ class OxoAggregatedKnowledgeBaseVulnerabilityType(graphene.ObjectType):
         description="List of vulnerabilities.",
     )
 
+    def resolve_highest_risk_rating(self, info) -> Optional[OxoRiskRatingEnum]:
+        try:
+            return OxoRiskRatingEnum[self.highest_risk_rating.name]
+        except KeyError:
+            return None
+
     def resolve_vulnerabilities(
         self: models.Scan,
         info: graphql_base.ResolveInfo,
@@ -668,13 +691,22 @@ class AgentsType(graphene.ObjectType):
     """Graphene object type for a list of agents."""
 
     agents = graphene.List(AgentType, required=True)
+    page_info = graphene.Field(
+        common.PageInfo,
+        required=False,
+    )
 
 
 class AgentGroupType(graphene_sqlalchemy.SQLAlchemyObjectType):
     """SQLAlchemy object type for an agent group."""
 
     key = graphene.String()
-    agents = graphene.Field(AgentsType, required=True)
+    agents = graphene.Field(
+        AgentsType,
+        required=True,
+        page=graphene.Int(required=False),
+        number_elements=graphene.Int(required=False),
+    )
 
     class Meta:
         """Meta class for the agent group object type."""
@@ -699,7 +731,10 @@ def resolve_key(self: models.AgentGroup, info: graphql_base.ResolveInfo) -> str:
         return f"agentgroup//{self.name}"
 
     def resolve_agents(
-        self: models.AgentGroup, info: graphql_base.ResolveInfo
+        self: models.AgentGroup,
+        info: graphql_base.ResolveInfo,
+        page: int = None,
+        number_elements: int = DEFAULT_NUMBER_ELEMENTS,
     ) -> AgentsType:
         """Resolve agents query.
         Args:
@@ -708,14 +743,28 @@ def resolve_agents(
         Returns:
             AgentsType: List of agents.
         """
+        if number_elements <= 0:
+            return AgentsType(agents=[])
+
         with models.Database() as session:
             agents = (
                 session.query(models.AgentGroup)
                 .filter(models.AgentGroup.id == self.id)
                 .first()
                 .agents
             )
-            return AgentsType(agents=agents)
+            if page is not None and number_elements > 0:
+                p = common.Paginator(agents, number_elements)
+                page = p.get_page(page)
+                page_info = common.PageInfo(
+                    count=p.count,
+                    num_pages=p.num_pages,
+                    has_next=page.has_next(),
+                    has_previous=page.has_previous(),
+                )
+                return OxoVulnerabilitiesType(agents=page, page_info=page_info)
+            else:
+                return AgentsType(agents=agents)
 
 
 class AgentGroupsType(graphene.ObjectType):