From 36590dea900a9d865a86a8cc026664d81398d7f2 Mon Sep 17 00:00:00 2001
From: Derek Fitchett <135860892+dfitchett@users.noreply.github.com>
Date: Fri, 18 Aug 2023 11:47:47 -0700
Subject: [PATCH] Domain EE: 1821: Fix issue with possible log injection
 (#1928)

---
 domain-ee/ee-max-cfi-app/src/python_src/api.py        |  7 ++++---
 .../ee-max-cfi-app/src/python_src/pydantic_models.py  |  6 +++---
 .../ee-max-cfi-app/src/python_src/util/sanitizer.py   |  2 ++
 domain-ee/ee-max-cfi-app/tests/test_api.py            | 11 ++++++++++-
 4 files changed, 19 insertions(+), 7 deletions(-)
 create mode 100644 domain-ee/ee-max-cfi-app/src/python_src/util/sanitizer.py

diff --git a/domain-ee/ee-max-cfi-app/src/python_src/api.py b/domain-ee/ee-max-cfi-app/src/python_src/api.py
index ca68a0634d..2e5b226f84 100644
--- a/domain-ee/ee-max-cfi-app/src/python_src/api.py
+++ b/domain-ee/ee-max-cfi-app/src/python_src/api.py
@@ -5,6 +5,7 @@
 from pydantic_models import (MaxRatingsForClaimForIncreaseRequest,
                              MaxRatingsForClaimForIncreaseResponse)
 from util.lookup_table import MAX_RATINGS_BY_CODE, get_max_rating
+from util.sanitizer import sanitize
 
 app = FastAPI(
     title="Max Ratings for CFI",
@@ -43,12 +44,12 @@ def get_health_status():
 def get_max_ratings(
         claim_for_increase: MaxRatingsForClaimForIncreaseRequest, ) -> MaxRatingsForClaimForIncreaseResponse:
     ratings = []
-    for dc in set(claim_for_increase.diagnostic_codes):
+    for dc in claim_for_increase.diagnostic_codes:
         validate_diagnostic_code(dc)
         max_rating = get_max_rating(dc)
         if max_rating:
             rating = {
-                "diagnostic_code": dc,
+                "diagnostic_code": sanitize(dc),
                 "max_rating": max_rating,
             }
             ratings.append(rating)
@@ -57,7 +58,7 @@ def get_max_ratings(
         "ratings": ratings
     }
 
-    logging.info(f"event=getMaxRating response={response}")
+    logging.info(f"event=getMaxRating ratings={ratings}")
     return response
 
 
diff --git a/domain-ee/ee-max-cfi-app/src/python_src/pydantic_models.py b/domain-ee/ee-max-cfi-app/src/python_src/pydantic_models.py
index ac4f21ad99..3c7c3f0fac 100644
--- a/domain-ee/ee-max-cfi-app/src/python_src/pydantic_models.py
+++ b/domain-ee/ee-max-cfi-app/src/python_src/pydantic_models.py
@@ -1,13 +1,13 @@
-from pydantic import BaseModel, conlist
+from pydantic import BaseModel, conint, conlist
 
 
 class MaxRatingsForClaimForIncreaseRequest(BaseModel):
-    diagnostic_codes: conlist(int, unique_items=True, min_items=1, max_items=100)
+    diagnostic_codes: conlist(conint(strict=True), unique_items=True, min_items=1, max_items=100)
 
 
 class Rating(BaseModel):
     diagnostic_code: int
-    max_rating: float
+    max_rating: int
 
 
 class MaxRatingsForClaimForIncreaseResponse(BaseModel):
diff --git a/domain-ee/ee-max-cfi-app/src/python_src/util/sanitizer.py b/domain-ee/ee-max-cfi-app/src/python_src/util/sanitizer.py
new file mode 100644
index 0000000000..cec818dea3
--- /dev/null
+++ b/domain-ee/ee-max-cfi-app/src/python_src/util/sanitizer.py
@@ -0,0 +1,2 @@
+def sanitize(obj) -> str:
+    return str(obj).replace('\r\n', '').replace('\n', '')
diff --git a/domain-ee/ee-max-cfi-app/tests/test_api.py b/domain-ee/ee-max-cfi-app/tests/test_api.py
index c3fd10e890..0b4e7fd7b5 100644
--- a/domain-ee/ee-max-cfi-app/tests/test_api.py
+++ b/domain-ee/ee-max-cfi-app/tests/test_api.py
@@ -94,10 +94,19 @@ def test_missing_params(client: TestClient):
     assert response.status_code == 422
 
 
-def test_unprocessable_content(client: TestClient):
+def test_unprocessable_content_request_does_not_have_array(client: TestClient):
     json_post_dict = {
         "diagnostic_codes": 6510,  # Should be an array
     }
 
     response = client.post(MAX_RATING, json=json_post_dict)
     assert response.status_code == 422
+
+
+def test_unprocessable_content_request_array_has_non_int_value(client: TestClient):
+    json_post_dict = {
+        "diagnostic_codes": ["6510"],  # Should be an int
+    }
+
+    response = client.post(MAX_RATING, json=json_post_dict)
+    assert response.status_code == 422