[NDR-253] Only search PDM table on requests from PDM (#812)

Co-authored-by: Megan <[email protected]>

Author: jameslinnell

lambdas/handlers/fhir_document_reference_search_handler.py

@@ -63,6 +63,7 @@ def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
         return_fhir=True,
         additional_filters=search_filters,
         check_upload_completed=False,
+        api_request_context=event.get("requestContext", {}),
     )
 
     if not document_references:

lambdas/services/document_reference_search_service.py

@@ -6,6 +6,7 @@
 from enums.dynamo_filter import AttributeOperator
 from enums.lambda_error import LambdaError
 from enums.metadata_field_names import DocumentReferenceMetadataFields
+from enums.mtls import MtlsCommonNames
 from enums.snomed_codes import SnomedCodes
 from models.document_reference import DocumentReference
 from models.fhir.R4.bundle import Bundle, BundleEntry
@@ -17,6 +18,7 @@
 from utils.dynamo_query_filter_builder import DynamoQueryFilterBuilder
 from utils.exceptions import DynamoServiceException
 from utils.lambda_exceptions import DocumentRefSearchException
+from utils.lambda_header_utils import validate_common_name_in_mtls
 
 logger = LoggingService(__name__)
 
@@ -28,6 +30,7 @@ def get_document_references(
         return_fhir: bool = False,
         additional_filters=None,
         check_upload_completed=True,
+        api_request_context: dict = {},
     ):
         """
         Fetch document references for a given NHS number.
@@ -38,8 +41,11 @@ def get_document_references(
         :param check_upload_completed: If True, check if the upload is completed before returning the results.
         :return: List of document references or FHIR DocumentReferences.
         """
+        common_name = validate_common_name_in_mtls(
+            api_request_context=api_request_context
+        )
         try:
-            list_of_table_names = self._get_table_names()
+            list_of_table_names = self._get_table_names(common_name)
             results = self._search_tables_for_documents(
                 nhs_number,
                 list_of_table_names,
@@ -60,13 +66,18 @@ def get_document_references(
             )
             raise DocumentRefSearchException(500, LambdaError.DocRefClient)
 
-    def _get_table_names(self) -> list[str]:
+    def _get_table_names(self, common_name: MtlsCommonNames | None) -> list[str]:
+        table_list = []
         try:
-            return json.loads(os.environ["DYNAMODB_TABLE_LIST"])
+            table_list = json.loads(os.environ["DYNAMODB_TABLE_LIST"])
         except JSONDecodeError as e:
             logger.error(f"Failed to decode table list: {str(e)}")
             raise
 
+        if not common_name or common_name not in MtlsCommonNames:
+            return table_list
+        return [t for t in table_list if common_name.value.lower() in t.lower()]
+
     def _search_tables_for_documents(
         self,
         nhs_number: str,
@@ -197,7 +208,7 @@ def create_document_reference_fhir_response(
             or document_reference.created,
             url=document_retrieve_endpoint
             + "/"
-            + SnomedCodes.LLOYD_GEORGE.value.code
+            + document_reference.document_snomed_code_type
             + "~"
             + document_reference.id,
         )

lambdas/services/post_fhir_document_reference_service.py

@@ -6,7 +6,6 @@
 from botocore.exceptions import ClientError
 from enums.lambda_error import LambdaError
 from enums.mtls import MtlsCommonNames
-from enums.patient_ods_inactive_status import PatientOdsInactiveStatus
 from enums.snomed_codes import SnomedCode, SnomedCodes
 from models.document_reference import DocumentReference
 from models.fhir.R4.fhir_document_reference import SNOMED_URL, Attachment
@@ -28,7 +27,6 @@
 )
 from utils.lambda_exceptions import CreateDocumentRefException
 from utils.lambda_header_utils import validate_common_name_in_mtls
-from utils.ods_utils import PCSE_ODS_CODE
 from utils.utilities import create_reference_id, get_pds_service, validate_nhs_number
 
 logger = LoggingService(__name__)

lambdas/tests/e2e/api/fhir/conftest.py

@@ -1,9 +1,11 @@
 import os
+import shutil
+import subprocess
+import tempfile
 import time
 
 import pytest
 import requests
-from syrupy.extensions.json import JSONSnapshotExtension
 
 from lambdas.tests.e2e.helpers.pdm_data_helper import PdmDataHelper
 
@@ -45,12 +47,39 @@ def fetch_with_retry_mtls(
     raise Exception("Condition not met within retry limit")
 
 
-@pytest.fixture
-def snapshot_json(snapshot):
-    return snapshot.with_defaults(extension_class=JSONSnapshotExtension)
-
-
 def create_mtls_session():
     session = requests.Session()
     session.cert = (CLIENT_CERT_PATH, CLIENT_KEY_PATH)
     return session
+
+
+@pytest.fixture
+def temp_cert_and_key():
+    import os
+
+    temp_dir = tempfile.mkdtemp()
+    key_path = os.path.join(temp_dir, "temp_key.pem")
+    cert_path = os.path.join(temp_dir, "temp_cert.pem")
+
+    try:
+        subprocess.run(["openssl", "genrsa", "-out", key_path, "2048"], check=True)
+        subprocess.run(
+            [
+                "openssl",
+                "req",
+                "-new",
+                "-x509",
+                "-key",
+                key_path,
+                "-out",
+                cert_path,
+                "-days",
+                "1",
+                "-subj",
+                "/C=GB/O=Test Org/CN=localhost",
+            ],
+            check=True,
+        )
+        yield cert_path, key_path
+    finally:
+        shutil.rmtree(temp_dir)

lambdas/tests/e2e/api/fhir/test_retrieve_document_fhir_api.py

@@ -1,3 +1,4 @@
+import base64
 import io
 import uuid
 
@@ -56,52 +57,68 @@ def test_file_retrieval(test_data, file_size):
     response = get_document_reference(pdm_record["id"])
     assert response.status_code == 200
 
-    if file_size:
-        json = response.json()
+    json = response.json()
+
+    if not file_size:
+        data = json["content"][0]["attachment"]["data"]
+        decoded_data = base64.b64decode(data)
+        expected_bytes = b"Sample PDF Content"
+        assert decoded_data == expected_bytes
+    else:
         expected_presign_uri = (
             f"https://{PDM_S3_BUCKET}.s3.eu-west-2.amazonaws.com/"
             f"{pdm_record['nhs_number']}/{pdm_record['id']}"
         )
         assert expected_presign_uri in json["content"][0]["attachment"]["url"]
 
 
-def test_no_file_found():
-    """Test retrieval when file does not exist."""
-    pdm_record = build_pdm_record()
-    response = get_document_reference(pdm_record["id"])
-    assert response.status_code == 404
+@pytest.mark.parametrize(
+    "nhs_id,expected_status,expected_code,expected_diagnostics",
+    [
+        ("9912003071", 404, "RESOURCE_NOT_FOUND", "Document reference not found"),
+    ],
+)
+def test_retrieve_edge_cases(
+    nhs_id, expected_status, expected_code, expected_diagnostics
+):
+    response = get_document_reference(nhs_id)
+    assert response.status_code == expected_status
+
+    body = response.json()
+    issue = body["issue"][0]
+
+    details = issue.get("details", {})
+    coding = details.get("coding", [{}])[0]
+
+    assert coding.get("code") == expected_code
+    assert issue.get("diagnostics") == expected_diagnostics
 
 
 def test_preliminary_file(test_data):
-    """Test retrieval for preliminary document status."""
     pdm_record = build_pdm_record(doc_status="preliminary")
     test_data.append(pdm_record)
-
     pdm_data_helper.create_metadata(pdm_record)
     pdm_data_helper.create_resource(pdm_record)
 
     response = get_document_reference(pdm_record["id"])
     assert response.status_code == 200
 
+    response_json = response.json()
+    assert response_json.get("docStatus") == "preliminary"
 
-# the error is about the auth token not mtls?
-def test_forbidden_without_mtls(test_data):
+
+def test_forbidden_with_invalid_cert(test_data, temp_cert_and_key):
     pdm_record = build_pdm_record()
     test_data.append(pdm_record)
-
     pdm_data_helper.create_metadata(pdm_record)
     pdm_data_helper.create_resource(pdm_record)
 
-    url = (
-        f"https://{MTLS_ENDPOINT}/FhirDocumentReference/{PDM_SNOMED}~{pdm_record['id']}"
-    )
-    headers = {
-        "Authorization": "Bearer 123",
-        "X-Correlation-Id": "1234",
-    }
-    response = requests.request("GET", url, headers=headers)
-    print("inital response:", response)
-    json = response.json()
-    print("response:", json)
+    # Use an invalid cert that is trusted by TLS but fails truststore validation
+    cert_path, key_path = temp_cert_and_key
+    url = f"https://{MTLS_ENDPOINT}/DocumentReference/{PDM_SNOMED}~{pdm_record['id']}"
+    headers = {"Authorization": "Bearer 123", "X-Correlation-Id": "1234"}
 
+    response = requests.get(url, headers=headers, cert=(cert_path, key_path))
+    body = response.json()
     assert response.status_code == 403
+    assert body["message"] == "Forbidden"

lambdas/tests/e2e/api/fhir/test_search_patient_fhir_api.py

@@ -1,10 +1,15 @@
 import io
-import logging
 import uuid
 
 import pytest
+import requests
 
-from lambdas.tests.e2e.api.fhir.conftest import MTLS_ENDPOINT, create_mtls_session
+from lambdas.tests.e2e.api.fhir.conftest import (
+    MTLS_ENDPOINT,
+    PDM_SNOMED,
+    create_mtls_session,
+)
+from lambdas.tests.e2e.conftest import APIM_ENDPOINT
 from lambdas.tests.e2e.helpers.pdm_data_helper import PdmDataHelper
 
 pdm_data_helper = PdmDataHelper()
@@ -46,33 +51,76 @@ def create_and_store_record(test_data, nhs_number="9912003071", doc_status=None)
 
 
 def test_search_patient_details(test_data):
-    """Search for a patient with one record."""
     create_and_store_record(test_data)
+
     response = search_document_reference("9912003071")
     assert response.status_code == 200
+
     bundle = response.json()
-    logging.info(bundle)
-    assert "entry" in bundle  # Basic validation
+    assert "entry" in bundle
+
+    attachment_url = bundle["entry"][0]["resource"]["content"][0]["attachment"]["url"]
+    assert (
+        f"https://{APIM_ENDPOINT}/national-document-repository/DocumentReference/{PDM_SNOMED}~"
+        in attachment_url
+    )
 
 
 def test_multiple_cancelled_search_patient_details(test_data):
-    """Search for a patient with multiple cancelled records."""
     create_and_store_record(test_data, doc_status="cancelled")
     create_and_store_record(test_data, doc_status="cancelled")
+
     response = search_document_reference("9912003071")
     assert response.status_code == 200
+
     bundle = response.json()
-    assert len(bundle.get("entry", [])) >= 2
+    entries = bundle.get("entry", [])
+    assert len(entries) >= 2
+
+    # Assert that all entries have status "cancelled"
+    for entry in entries:
+        resource = entry.get("resource", {})
+        assert resource.get("docStatus") == "cancelled"
 
 
 @pytest.mark.parametrize(
-    "nhs_number,expected_status",
+    "nhs_number,expected_status,expected_code,expected_diagnostics",
     [
-        ("9912003071", 404),  # No records
-        ("9999999993", 400),  # Invalid patient
+        ("9912003071", 404, "RESOURCE_NOT_FOUND", "Document reference not found"),
+        ("9999999993", 400, "INVALID_SEARCH_DATA", "Invalid patient number 9999999993"),
+        ("123", 400, "INVALID_SEARCH_DATA", "Invalid patient number 123"),
     ],
 )
-def test_search_edge_cases(nhs_number, expected_status):
-    """Test search for no records and invalid patient."""
+def test_search_edge_cases(
+    nhs_number, expected_status, expected_code, expected_diagnostics
+):
     response = search_document_reference(nhs_number)
     assert response.status_code == expected_status
+
+    body = response.json()
+    issue = body["issue"][0]
+    details = issue.get("details", {})
+    coding = details.get("coding", [{}])[0]
+    assert coding.get("code") == expected_code
+    assert issue.get("diagnostics") == expected_diagnostics
+
+
+def test_search_patient_unauthorized_mtls(test_data, temp_cert_and_key):
+    """Search should return 403 when mTLS certificate is invalid or missing."""
+    create_and_store_record(test_data)
+    url = (
+        f"https://{MTLS_ENDPOINT}/DocumentReference?"
+        f"subject:identifier=https://fhir.nhs.uk/Id/nhs-number|9912003071"
+    )
+    headers = {
+        "Authorization": "Bearer 123",
+        "X-Correlation-Id": "unauthorized-test",
+    }
+
+    # Use an invalid cert that is trusted by TLS but fails truststore validation
+    cert_path, key_path = temp_cert_and_key
+
+    response = requests.get(url, headers=headers, cert=(cert_path, key_path))
+    body = response.json()
+    assert response.status_code == 403
+    assert body["message"] == "Forbidden"

lambdas/tests/unit/conftest.py

@@ -149,7 +149,7 @@ def set_env(monkeypatch):
     monkeypatch.setenv(MOCK_PDM_TABLE_NAME_ENV_NAME, MOCK_PDM_TABLE_NAME)
     monkeypatch.setenv(MOCK_PDM_BUCKET_ENV_NAME, MOCK_PDM_BUCKET)
     monkeypatch.setenv(
-        "DYNAMODB_TABLE_LIST", json.dumps([MOCK_ARF_TABLE_NAME, MOCK_LG_TABLE_NAME])
+        "DYNAMODB_TABLE_LIST", json.dumps([MOCK_PDM_TABLE_NAME, MOCK_LG_TABLE_NAME])
     )
     monkeypatch.setenv(MOCK_ZIP_OUTPUT_BUCKET_ENV_NAME, MOCK_ZIP_OUTPUT_BUCKET)
     monkeypatch.setenv(MOCK_ZIP_TRACE_TABLE_ENV_NAME, MOCK_ZIP_TRACE_TABLE)

lambdas/tests/unit/handlers/test_fhir_document_reference_search_handler.py

@@ -147,6 +147,7 @@ def test_lambda_handler_returns_200_with_documents(
         return_fhir=True,
         additional_filters={},
         check_upload_completed=False,
+        api_request_context={},
     )
 
 
@@ -163,6 +164,7 @@ def test_lambda_handler_returns_404_when_no_documents(
         return_fhir=True,
         additional_filters={},
         check_upload_completed=False,
+        api_request_context={},
     )
 
 
@@ -208,6 +210,7 @@ def test_lambda_handler_with_additional_filters(
         return_fhir=True,
         additional_filters=expected_filters,
         check_upload_completed=False,
+        api_request_context={},
     )