Skip to content
Merged
Show file tree
Hide file tree
Changes from 12 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
89 changes: 89 additions & 0 deletions usaspending_api/common/api_request_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import boto3
from botocore.exceptions import ClientError
from django.contrib.postgres.search import SearchVector
from django.core.exceptions import FieldDoesNotExist
from django.db.models import Q
from django.utils import timezone
from rest_framework import status
Expand Down Expand Up @@ -118,6 +119,13 @@ def create_from_query_params(self, parameters: Dict[str, Any]) -> Dict[str, Any]
for key in parameters:
if key in self.ignored_parameters:
continue

# Determine the actual field that will be used after mapping.
actual_field = self.filter_map.get(key, key)

# Validate the field path to prevent ORM injection.
self._resolve_field_path(actual_field)

if key in self.filter_map:
return_arguments[self.filter_map[key]] = parameters[key]
else:
Expand Down Expand Up @@ -260,10 +268,91 @@ def _validate_single_filter(self, filt: Dict[str, Any]) -> None:
if "combine_method" in filt:
self.validate_post_request(filt)
elif "field" in filt and "operation" in filt and "value" in filt:
self._validate_filter_field(filt)
self._validate_filter_operation(filt)
else:
raise InvalidParameterException("Malformed filter - missing field, operation, or value")

def _validate_filter_field(self, filt: Dict[str, Any]) -> None:
"""Validate the field parameter to prevent ORM injection."""
field = filt["field"]

# Special cases: field arrays for range_intersect and search operations.
if isinstance(field, list):
operation = filt.get("operation", "")
if operation.replace("not_", "") in ["range_intersect", "search"]:
# Validate each field in the array.
for single_field in field:
self._resolve_field_path(single_field)
else:
raise InvalidParameterException("Field arrays are only allowed for range_intersect and search ops.")
else:
# Standard case: validate single field.
self._resolve_field_path(field)

def _resolve_field_path(self, field_path: str) -> None:
"""Validate field path against model schema, allowing FK traversal and known operations."""
segments = field_path.split("__")
current_model = self.model
resolved_field = None

for index, segment in enumerate(segments):
try:
resolved_field = current_model._meta.get_field(segment)
except FieldDoesNotExist:
# If this is the last segment, check if it's a known Django lookup/operation.
is_last = index == len(segments) - 1
if is_last and segment in self._get_allowed_lookups():
# Valid lookup suffix, no need to continue validation.
return
raise InvalidParameterException(f"Invalid field: {field_path}") from None

is_last = index == len(segments) - 1
if resolved_field.is_relation and not is_last:
current_model = resolved_field.related_model
elif not is_last:
# Non-relation field with more segments. Check if remaining segments are all valid lookups.
remaining_segments = segments[index + 1:]
allowed_lookups = self._get_allowed_lookups()
if all(seg in allowed_lookups for seg in remaining_segments):
# All remaining segments are valid lookups, validation complete.
return
raise InvalidParameterException(f"Invalid field: {field_path}")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Now uses f"" strings throughout per suggestion.


def _get_allowed_lookups(self) -> set:
"""Return set of allowed Django field lookups that can appear as a final segment."""
# Derive allowed lookups from operators dictionary to reduce duplication.
allowed = set()

for operation_value in self.operators.values():
if operation_value.startswith("__"):
# Extract lookup suffix: "__lt" → "lt", "__icontains" → "icontains".
lookup = operation_value[2:]

# Handle compound lookups like "__len__gt" → ["len", "gt"].
# These represent transform + lookup chains (e.g., field__len__gt).
if "__" in lookup:
allowed.update(lookup.split("__"))
else:
allowed.add(lookup)
elif operation_value == "":
# "equals" operation maps to "" which becomes "exact" lookup.
allowed.add("exact")

# Add additional safe Django lookups not covered by operators.
allowed.update({
"iexact", "startswith", "istartswith", "endswith", "iendswith",
"date", "year", "month", "day", "week", "week_day", "quarter",
"time", "hour", "minute", "second"
})

# Explicitly exclude dangerous lookups to prevent ReDoS attacks.
# NOTE: These must never be allowed even if accidentally added to operators dict.
allowed.discard("regex")
allowed.discard("iregex")

return allowed

def _validate_filter_operation(self, filt: Dict[str, Any]) -> None:
"""Validate the operation and value for a filter"""
operation = filt["operation"]
Expand Down
63 changes: 53 additions & 10 deletions usaspending_api/common/mixins.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,12 @@
from django.db.models import Avg, Count, F, Q, Max, Min, Sum, Func, IntegerField, ExpressionWrapper
from typing import Any

from django.core.exceptions import FieldDoesNotExist
from django.db.models import Avg, Count, ExpressionWrapper, F, Func, IntegerField, Max, Min, Model, Q, QuerySet, Sum
from django.db.models.fields import Field
from django.db.models.functions import ExtractDay, ExtractMonth, ExtractYear
from rest_framework.request import Request

from usaspending_api.common.api_request_utils import FilterGenerator, AutoCompleteHandler
from usaspending_api.common.api_request_utils import AutoCompleteHandler, FilterGenerator
from usaspending_api.common.exceptions import InvalidParameterException


Expand All @@ -12,7 +17,7 @@ class AggregateQuerysetMixin(object):
in the get_queryset method).
"""

def aggregate(self, request, *args, **kwargs):
def aggregate(self, request: Request, *args, **kwargs) -> QuerySet:
"""Perform an aggregate function on a Django queryset with an optional group by field."""
# create a single dict that contains the requested aggregate parameters, regardless of request type
# (e.g., GET, POST) (not sure if this is a good practice, or we should be more prescriptive that aggregate
Expand Down Expand Up @@ -72,7 +77,41 @@ def aggregate(self, request, *args, **kwargs):

_sql_function_transformations = {"fy": IntegerField}

def _wrapped_f_expression(self, col_name):
def _resolve_field_path(self, model: type[Model], field_path: str) -> tuple[str, Field | None]:
segments = field_path.split("__")
current_model = model
resolved_field = None

for index, segment in enumerate(segments):
try:
resolved_field = current_model._meta.get_field(segment)
except FieldDoesNotExist:
# If this is the last segment and it's a known transform, allow it.
is_last = index == len(segments) - 1
if is_last and segment in self._sql_function_transformations:
# Valid transform suffix, no need to continue validation.
return field_path, None
raise InvalidParameterException(f"Invalid field: {field_path}") from None

is_last = index == len(segments) - 1
if resolved_field.is_relation and not is_last:
current_model = resolved_field.related_model
elif not is_last:
# Check if the next segment is a valid SQL function transformation
next_segment = segments[index + 1]
is_next_last = (index + 1) == len(segments) - 1
if is_next_last and next_segment in self._sql_function_transformations:
# Valid transform suffix on this field, allow it
return field_path, None
raise InvalidParameterException(f"Invalid field: {field_path}")

return field_path, resolved_field

def _validate_field_path(self, model: type[Model], field_path: str) -> str:
self._resolve_field_path(model, field_path)
return field_path

def _wrapped_f_expression(self, col_name: str) -> F | ExpressionWrapper:
"""F-expression of col, wrapped if needed with SQL function call

Assumes that there's an SQL function defined for each registered lookup.
Expand All @@ -87,7 +126,7 @@ def _wrapped_f_expression(self, col_name):
return result
return F(col_name)

def validate_request(self, params, queryset):
def validate_request(self, params: dict[str, Any], queryset: QuerySet) -> tuple[str, list[str], str | None]:
"""Validate request parameters."""

agg_field = params.get("field")
Expand Down Expand Up @@ -131,6 +170,9 @@ def validate_request(self, params, queryset):
if not isinstance(group_fields, list):
group_fields = [group_fields]

for group_field in group_fields:
self._validate_field_path(model, group_field)

# if a groupby date part is specified, make sure the groupby field is
# a date and the groupby value is year, quarter, or month
if date_part is not None:
Expand All @@ -140,7 +182,8 @@ def validate_request(self, params, queryset):
# if the request is asking to group by a date component, the field
# we're grouping by must be a date-related field (there is probably a better way to do this?)
date_fields = ["DateField", "DateTimeField"]
if model._meta.get_field(group_fields[0]).get_internal_type() not in date_fields:
_, group_field = self._resolve_field_path(model, group_fields[0])
if group_field.get_internal_type() not in date_fields:
raise InvalidParameterException(
"Group by date part ({}) requested for a non-date group by ({})".format(date_part, group_fields[0])
)
Expand All @@ -158,7 +201,7 @@ def validate_request(self, params, queryset):
class FilterQuerysetMixin(object):
"""Handles queryset filtering."""

def filter_records(self, request, *args, **kwargs):
def filter_records(self, request: Request, *args, **kwargs) -> QuerySet:
"""Filter a queryset based on request parameters"""
queryset = kwargs.get("queryset")

Expand Down Expand Up @@ -195,7 +238,7 @@ def filter_records(self, request, *args, **kwargs):

return queryset.filter(subwhere)

def order_records(self, request, *args, **kwargs):
def order_records(self, request: Request, *args, **kwargs) -> QuerySet:
"""Order a queryset based on request parameters."""
queryset = kwargs.get("queryset")

Expand All @@ -212,7 +255,7 @@ def order_records(self, request, *args, **kwargs):
else:
return queryset

def get_submission_id_filters(self):
def get_submission_id_filters(self) -> tuple[int | None, list[int]]:
"""
Returns the federal_account_id and the list of fiscal_years from the list of incoming
filters if they exist. If not, return None and an empty list respectively
Expand All @@ -231,7 +274,7 @@ def get_submission_id_filters(self):
class AutocompleteResponseMixin(object):
"""Handles autocomplete responses and requests"""

def build_response(self, request, *args, **kwargs):
def build_response(self, request: Request, *args, **kwargs) -> Any:
queryset = kwargs.get("queryset")

serializer = kwargs.get("serializer")
Expand Down
62 changes: 59 additions & 3 deletions usaspending_api/common/tests/integration/test_aggregate_mixin.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,10 @@
from unittest.mock import Mock

import pytest
from model_bakery.recipe import Recipe
from model_bakery import baker
from model_bakery.recipe import Recipe

from usaspending_api.awards.models import Award
from usaspending_api.awards.models import TransactionNormalized
from usaspending_api.awards.models import Award, TransactionNormalized
from usaspending_api.common.mixins import AggregateQuerysetMixin
from usaspending_api.search.models import AwardSearch

Expand Down Expand Up @@ -353,3 +352,60 @@ def itemsorter(a):
assert agg_list[0]["aggregate"] is None
assert agg_list[1]["aggregate"] == 10.0
assert agg_list[2]["aggregate"] == 30.0


@pytest.mark.django_db
def test_aggregate_group_with_fk_traversal(client, aggregate_models):
"""Test that legitimate FK traversal in group fields works."""
response = client.post(
"/api/v1/awards/total/",
{"field": "total_obligation", "aggregate": "sum", "group": "type"},
content_type="application/json",
)
assert response.status_code == 200


@pytest.mark.django_db
def test_aggregate_group_with_fy_transform(client, aggregate_models):
"""Test that __fy transform suffix works."""
response = client.post(
"/api/v1/awards/total/",
{"field": "total_obligation", "aggregate": "sum", "group": "period_of_performance_start_date__fy"},
content_type="application/json",
)
assert response.status_code == 200


@pytest.mark.django_db
def test_aggregate_group_blocks_regex_injection(client, aggregate_models):
"""Test that __regex injection is blocked in aggregate group fields."""
response = client.post(
"/api/v1/awards/total/",
{"field": "total_obligation", "aggregate": "sum", "group": "type__regex"},
content_type="application/json",
)
assert response.status_code == 400
assert "Invalid field" in response.json()["detail"]


@pytest.mark.django_db
def test_aggregate_group_blocks_iregex_injection(client, aggregate_models):
"""Test that __iregex injection is blocked in aggregate group fields."""
response = client.post(
"/api/v1/awards/total/",
{"field": "total_obligation", "aggregate": "sum", "group": "type__iregex"},
content_type="application/json",
)
assert response.status_code == 400
assert "Invalid field" in response.json()["detail"]


@pytest.mark.django_db
def test_aggregate_group_blocks_invalid_field_path(client, aggregate_models):
"""Test that non-existent fields in group are rejected."""
response = client.post(
"/api/v1/awards/total/",
{"field": "total_obligation", "aggregate": "sum", "group": "nonexistent_field__name"},
content_type="application/json",
)
assert response.status_code == 400
Loading
Loading