Skip to content

Commit

Permalink
Merge branch 'main' into fix/delete_agent_group_mapping_on_deletion_m…
Browse files Browse the repository at this point in the history 
…utation
  • Loading branch information
@adnaneserrar
adnaneserrar authored Aug 7, 2024
2 parents 9f7b4b0 + 79e56da commit 2bf8a2d
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 51 deletions.
62 changes: 35 additions & 27 deletions src/ostorlab/serve_app/oxo.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -292,25 +292,25 @@ def mutate(
return ExportScanMutation(content=export_file_content)


class DeleteScanMutation(graphene.Mutation):
class DeleteScansMutation(graphene.Mutation):
"""Delete Scan & its information mutation."""

class Arguments:
scan_id = graphene.Int(required=True)
scan_ids = graphene.List(graphene.Int, required=True)

result = graphene.Boolean()

@staticmethod
def mutate(
root,
info: graphql_base.ResolveInfo,
scan_id: int,
) -> "DeleteScanMutation":
scan_ids: list[int],
) -> "DeleteScansMutation":
"""Delete a scan & its information.
Args:
info: `graphql_base.ResolveInfo` instance.
scan_id: The scan ID.
scan_ids: The scan IDs.
Raises:
graphql.GraphQLError in case the scan does not exist.
Expand All @@ -320,18 +320,23 @@ def mutate(
"""
with models.Database() as session:
scan_query = session.query(models.Scan).filter_by(id=scan_id)
if scan_query.count() == 0:
raise graphql.GraphQLError("Scan not found.")
scan_query.delete()
session.query(models.Vulnerability).filter_by(scan_id=scan_id).delete()
session.query(models.ScanStatus).filter_by(scan_id=scan_id).delete()
DeleteScanMutation._delete_assets(scan_id, session)
session.commit()
return DeleteScanMutation(result=True)
scans = (
session.query(models.Scan).filter(models.Scan.id.in_(scan_ids)).all()
)
if len(scans) == 0:
raise graphql.GraphQLError("No scan is found.")

for scan in scans:
scan_query = session.query(models.Scan).filter_by(id=scan.id)
scan_query.delete()
session.query(models.Vulnerability).filter_by(scan_id=scan.id).delete()
session.query(models.ScanStatus).filter_by(scan_id=scan.id).delete()
DeleteScansMutation.delete_assets(scan.id, session)
session.commit()
return DeleteScansMutation(result=True)

@staticmethod
def _delete_assets(scan_id: int, session: models.Database) -> None:
def delete_assets(scan_id: int, session: models.Database) -> None:
"""Delete assets.
Args:
Expand Down Expand Up @@ -481,21 +486,21 @@ def _validate(asset: types.OxoAssetInputType) -> Optional[str]:
return None


class StopScanMutation(graphene.Mutation):
class StopScansMutation(graphene.Mutation):
"""Stop scan mutation."""

class Arguments:
scan_id = graphene.Int(required=True)
scan_ids = graphene.List(graphene.Int, required=True)

scan = graphene.Field(types.OxoScanType)
scans = graphene.List(types.OxoScanType)

@staticmethod
def mutate(root, info: graphql_base.ResolveInfo, scan_id: int):
def mutate(root, info: graphql_base.ResolveInfo, scan_ids: list[int]):
"""Stop the desired scan.
Args:
info: `graphql_base.ResolveInfo` instance.
scan_id: The scan ID.
scan_ids: The scan IDs.
Raises:
graphql.GraphQLError in case the scan does not exist or the scan id is invalid.
Expand All @@ -505,11 +510,14 @@ def mutate(root, info: graphql_base.ResolveInfo, scan_id: int):
"""
with models.Database() as session:
scan = session.query(models.Scan).get(scan_id)
if scan is None:
raise graphql.GraphQLError("Scan not found.")
local_runtime.LocalRuntime().stop(scan_id=str(scan_id))
return StopScanMutation(scan=scan)
scans = (
session.query(models.Scan).filter(models.Scan.id.in_(scan_ids)).all()
)
if len(scans) == 0:
raise graphql.GraphQLError("No scan is found.")
for scan_id in scan_ids:
local_runtime.LocalRuntime().stop(scan_id=str(scan_id))
return StopScansMutation(scans=scans)


class PublishAgentGroupMutation(graphene.Mutation):
Expand Down Expand Up @@ -874,7 +882,7 @@ def _run_scan_background(


class Mutations(graphene.ObjectType):
delete_scan = DeleteScanMutation.Field(
delete_scans = DeleteScansMutation.Field(
description="Delete a scan & all its information."
)
delete_agent_group = DeleteAgentGroupMutation.Field(
Expand All @@ -883,7 +891,7 @@ class Mutations(graphene.ObjectType):
import_scan = ImportScanMutation.Field(description="Import scan from file.")
export_scan = ExportScanMutation.Field(description="Export scan to file.")
create_assets = CreateAssetsMutation.Field(description="Create an asset.")
stop_scan = StopScanMutation.Field(
stop_scans = StopScansMutation.Field(
description="Stops running scan, scan is marked as stopped once the engine has completed cancellation."
)
publish_agent_group = PublishAgentGroupMutation.Field(
Expand Down
47 changes: 23 additions & 24 deletions tests/serve_app/oxo_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -707,7 +707,7 @@ def testQueryScan_whenScanDoesNotExist_returnErrorMessage(
assert response.get_json()["errors"][0]["message"] == "Scan not found."


def testDeleteScanMutation_whenScanExist_deleteScanAndVulnz(
def testDeleteScansMutation_whenScanExist_deleteScanAndVulnz(
authenticated_flask_client: testing.FlaskClient, android_scan: models.Scan
) -> None:
"""Ensure the delete scan mutation deletes the scan, its statuses & vulnerabilities."""
Expand All @@ -726,15 +726,15 @@ def testDeleteScanMutation_whenScanExist_deleteScanAndVulnz(
)

query = """
mutation DeleteScan ($scanId: Int!){
deleteScan (scanId: $scanId) {
mutation DeleteScans ($scanIds: [Int]!){
deleteScans (scanIds: $scanIds) {
result
}
}
"""

response = authenticated_flask_client.post(
"/graphql", json={"query": query, "variables": {"scanId": android_scan.id}}
"/graphql", json={"query": query, "variables": {"scanIds": [android_scan.id]}}
)

assert response.status_code == 200, response.get_json()
Expand All @@ -759,24 +759,24 @@ def testDeleteScanMutation_whenScanExist_deleteScanAndVulnz(
)


def testDeleteScanMutation_whenScanDoesNotExist_returnErrorMessage(
def testDeleteScansMutation_whenScanDoesNotExist_returnErrorMessage(
authenticated_flask_client: testing.FlaskClient,
) -> None:
"""Ensure the delete scan mutation returns an error message when the scan does not exist."""
query = """
mutation DeleteScan ($scanId: Int!){
deleteScan (scanId: $scanId) {
mutation DeleteScans ($scanIds: [Int]!){
deleteScans (scanIds: $scanIds) {
result
}
}
"""

response = authenticated_flask_client.post(
"/graphql", json={"query": query, "variables": {"scanId": 42}}
"/graphql", json={"query": query, "variables": {"scanIds": [42]}}
)

assert response.status_code == 200, response.get_json()
assert response.get_json()["errors"][0]["message"] == "Scan not found."
assert response.get_json()["errors"][0]["message"] == "No scan is found."


def testScansQuery_withPagination_shouldReturnPageInfo(
Expand Down Expand Up @@ -2012,13 +2012,13 @@ def testQueryAssets_whenScanHasMultipleAssets_shouldReturnAllAssets(
]


def testStopScanMutation_whenScanIsRunning_shouldStopScan(
def testStopScansMutation_whenScanIsRunning_shouldStopScan(
authenticated_flask_client: testing.FlaskClient,
in_progress_web_scan: models.Scan,
mocker: plugin.MockerFixture,
db_engine_path: str,
) -> None:
"""Test stopScan mutation when scan is running should stop scan."""
"""Test stopScans mutation when scan is running should stop scan."""
mocker.patch(
"ostorlab.cli.docker_requirements_checker.is_docker_installed",
return_value=True,
Expand Down Expand Up @@ -2047,56 +2047,55 @@ def testStopScanMutation_whenScanIsRunning_shouldStopScan(
scan = session.query(models.Scan).get(in_progress_web_scan.id)
scan_progress = scan.progress
query = """
mutation stopScan($scanId: Int!){
stopScan(scanId: $scanId){
scan{
mutation stopScans($scanIds: [Int]!){
stopScans(scanIds: $scanIds){
scans{
id
}
}
}
"""
response = authenticated_flask_client.post(
"/graphql", json={"query": query, "variables": {"scanId": str(scan.id)}}
"/graphql", json={"query": query, "variables": {"scanIds": [scan.id]}}
)

assert response.status_code == 200, response.get_json()
session.refresh(scan)
scan = session.query(models.Scan).get(in_progress_web_scan.id)
response_json = response.get_json()
nbr_scans_after = session.query(models.Scan).count()
assert response_json["data"] == {
"stopScan": {"scan": {"id": str(in_progress_web_scan.id)}}
"stopScans": {"scans": [{"id": str(in_progress_web_scan.id)}]}
}
assert scan.progress.name == "STOPPED"
assert scan.progress != scan_progress
assert nbr_scans_after == nbr_scans_before


def testStopScanMutation_whenNoScanFound_shouldReturnError(
def testStopScansMutation_whenNoScanFound_shouldReturnError(
authenticated_flask_client: testing.FlaskClient,
mocker: plugin.MockerFixture,
db_engine_path: str,
clean_db: None,
) -> None:
"""Test stopScan mutation when scan doesn't exist should return error message."""
"""Test stopScans mutation when scan doesn't exist should return error message."""
del clean_db
mocker.patch.object(models, "ENGINE_URL", db_engine_path)
query = """
mutation stopScan($scanId: Int!){
stopScan(scanId: $scanId){
scan{
mutation stopScans($scanIds: [Int]!){
stopScans(scanIds: $scanIds){
scans{
id
}
}
}
"""
response = authenticated_flask_client.post(
"/graphql", json={"query": query, "variables": {"scanId": "5"}}
"/graphql", json={"query": query, "variables": {"scanIds": ["5"]}}
)

assert response.status_code == 200, response.get_json()
response_json = response.get_json()
assert response_json["errors"][0]["message"] == "Scan not found."
assert response_json["errors"][0]["message"] == "No scan is found."


def testQueryVulnerabilitiesOfKb_withPagination_shouldReturnPageInfo(
Expand Down

0 comments on commit 2bf8a2d

Please sign in to comment.