changed api for fetching unmasked email

Dockerfile

@@ -15,9 +15,12 @@ COPY . .
 ENV PORT=7860
 ENV MODEL_PATH="Sparkonix/email-classifier-model" 
 
-# Change this to point to SQLite database location
+# SQLite database path
 ENV DATABASE_PATH="/data/emails.db"
 
+# Global access key for email retrieval
+ENV EMAIL_ACCESS_KEY="access_key_123"
+
 # Add this line to set cache location to a writable directory
 ENV HF_HOME="/app/.cache/huggingface"
 

Email_Classification_API_Tests.postman_collection.json

@@ -12,20 +12,14 @@
         "type": "string"
       },
       {
-        "key": "emailId",
-        "value": ""
+        "key": "maskedEmailWithPII",
+        "value": "",
+        "type": "string"
       },
       {
         "key": "accessKey",
-        "value": ""
-      },
-      {
-        "key": "piiEmailId",
-        "value": ""
-      },
-      {
-        "key": "piiAccessKey",
-        "value": ""
+        "value": "local_dev_secure_access_key_20240516",
+        "type": "string"
       }
     ],
     "item": [
@@ -125,15 +119,7 @@
                 "pm.test(\"Category is valid\", function() {",
                 "    var jsonData = pm.response.json();",
                 "    pm.expect([\"Incident\", \"Request\", \"Change\", \"Problem\"]).to.include(jsonData.category_of_the_email);",
-                "});",
-                "",
-                "// Store the email_id and access_key for database retrieval tests",
-                "var responseJson = pm.response.json();",
-                "if (responseJson.email_id && responseJson.access_key) {",
-                "    pm.collectionVariables.set(\"emailId\", responseJson.email_id);",
-                "    pm.collectionVariables.set(\"accessKey\", responseJson.access_key);",
-                "    console.log(\"Stored email_id and access_key for retrieval tests\");",
-                "}"
+                "});"
               ],
               "type": "text/javascript"
             }
@@ -191,13 +177,10 @@
                 "    pm.expect(jsonData.list_of_masked_entities).to.be.an(\"array\").that.is.not.empty;",
                 "});",
                 "",
-                "// Store the email_id and access_key for database retrieval tests specifically for PII emails",
+                "// Store the masked email for later unmasking test",
                 "var responseJson = pm.response.json();",
-                "if (responseJson.email_id && responseJson.access_key) {",
-                "    pm.collectionVariables.set(\"piiEmailId\", responseJson.email_id);",
-                "    pm.collectionVariables.set(\"piiAccessKey\", responseJson.access_key);",
-                "    console.log(\"Stored PII email_id and access_key for retrieval tests\");",
-                "}"
+                "pm.collectionVariables.set(\"maskedEmailWithPII\", responseJson.masked_email);",
+                "console.log(\"Stored masked email for unmasking tests\");"
               ],
               "type": "text/javascript"
             }
@@ -450,13 +433,13 @@
         "response": []
       },
       {
-        "name": "Retrieve Original Email",
+        "name": "Unmask Email by Content",
         "request": {
           "method": "POST",
           "header": [],
           "body": {
             "mode": "raw",
-            "raw": "{\n    \"email_id\": \"{{piiEmailId}}\",\n    \"access_key\": \"{{piiAccessKey}}\"\n}",
+            "raw": "{\n    \"masked_email\": \"{{maskedEmailWithPII}}\",\n    \"access_key\": \"{{accessKey}}\"\n}",
             "options": {
               "raw": {
                 "language": "json"
@@ -464,18 +447,17 @@
             }
           },
           "url": {
-            "raw": "{{baseUrl}}/api/v1/original-email/retrieve",
+            "raw": "{{baseUrl}}/api/v1/unmask-email",
             "host": [
               "{{baseUrl}}"
             ],
             "path": [
               "api",
               "v1",
-              "original-email",
-              "retrieve"
+              "unmask-email"
             ]
           },
-          "description": "Retrieve the original unmasked email from the database"
+          "description": "Retrieve the original unmasked email using the masked content"
         },
         "event": [
           {
@@ -503,8 +485,9 @@
                 "    pm.expect(jsonData.data).to.have.property(\"masked_email\");",
                 "    pm.expect(jsonData.data).to.have.property(\"masked_entities\");",
                 "    ",
-                "    // Check that the original email contains the PII (for the PII test email)",
-                "    if (pm.collectionVariables.get(\"piiEmailId\")) {",
+                "    // Check that the original email contains the PII",
+                "    var maskedEmailWithPII = pm.collectionVariables.get(\"maskedEmailWithPII\");",
+                "    if (maskedEmailWithPII) {",
                 "        pm.expect(jsonData.data.original_email).to.include(\"John Smith\");",
                 "        pm.expect(jsonData.data.original_email).to.include(\"john.smith@example.com\");",
                 "        pm.expect(jsonData.data.original_email).to.include(\"555-123-4567\");",
@@ -518,13 +501,13 @@
         "response": []
       },
       {
-        "name": "Retrieve With Invalid Access Key",
+        "name": "Unmask Email with Invalid Access Key",
         "request": {
           "method": "POST",
           "header": [],
           "body": {
             "mode": "raw",
-            "raw": "{\n    \"email_id\": \"{{piiEmailId}}\",\n    \"access_key\": \"invalid_access_key_123456\"\n}",
+            "raw": "{\n    \"masked_email\": \"{{maskedEmailWithPII}}\",\n    \"access_key\": \"invalid_access_key_123456\"\n}",
             "options": {
               "raw": {
                 "language": "json"
@@ -532,26 +515,75 @@
             }
           },
           "url": {
-            "raw": "{{baseUrl}}/api/v1/original-email/retrieve",
+            "raw": "{{baseUrl}}/api/v1/unmask-email",
             "host": [
               "{{baseUrl}}"
             ],
             "path": [
               "api",
               "v1",
-              "original-email",
-              "retrieve"
+              "unmask-email"
             ]
           },
-          "description": "Test security by attempting to retrieve email with invalid access key"
+          "description": "Test security by attempting to unmask email with invalid access key"
         },
         "event": [
           {
             "listen": "test",
             "script": {
               "exec": [
-                "// Check that we get an error (404) for invalid access key",
+                "// Check that we get an error (401) for invalid access key",
                 "pm.test(\"Should return error for invalid access key\", function() {",
+                "    pm.expect(pm.response.code).to.equal(401);",
+                "});",
+                "",
+                "// Check error message",
+                "pm.test(\"Response contains appropriate error message\", function() {",
+                "    var jsonData = pm.response.json();",
+                "    pm.expect(jsonData).to.have.property(\"detail\");",
+                "    pm.expect(jsonData.detail).to.include(\"Invalid access key\");",
+                "});"
+              ],
+              "type": "text/javascript"
+            }
+          }
+        ],
+        "response": []
+      },
+      {
+        "name": "Unmask with Non-existent Email Content",
+        "request": {
+          "method": "POST",
+          "header": [],
+          "body": {
+            "mode": "raw",
+            "raw": "{\n    \"masked_email\": \"This is a masked email that does not exist in the database [FULL_NAME].\",\n    \"access_key\": \"{{accessKey}}\"\n}",
+            "options": {
+              "raw": {
+                "language": "json"
+              }
+            }
+          },
+          "url": {
+            "raw": "{{baseUrl}}/api/v1/unmask-email",
+            "host": [
+              "{{baseUrl}}"
+            ],
+            "path": [
+              "api",
+              "v1",
+              "unmask-email"
+            ]
+          },
+          "description": "Test error handling when masked email content doesn't exist"
+        },
+        "event": [
+          {
+            "listen": "test",
+            "script": {
+              "exec": [
+                "// Check that we get a 404 error for non-existent email content",
+                "pm.test(\"Should return 404 for non-existent email content\", function() {",
                 "    pm.expect(pm.response.code).to.equal(404);",
                 "});",
                 "",

database.py

@@ -7,7 +7,6 @@ import sqlite3
 from typing import Dict, Any, Optional, List, Tuple
 from datetime import datetime
 import uuid
-import hashlib
 
 
 class EmailDatabase:
@@ -30,6 +29,9 @@ class EmailDatabase:
             "/data/emails.db"  # This path persists in Hugging Face Spaces
         )
         
+        # Get the global access key from environment variables
+        self.access_key = os.environ.get("EMAIL_ACCESS_KEY", "default_secure_access_key")
+        
         # Ensure the data directory exists
         self._ensure_data_directory()
         
@@ -64,14 +66,10 @@ class EmailDatabase:
                 masked_email TEXT NOT NULL,
                 masked_entities TEXT NOT NULL, 
                 category TEXT,
-                created_at TEXT NOT NULL,
-                access_key TEXT NOT NULL
+                created_at TEXT NOT NULL
             )
             ''')
             
-            # Create an index on the access_key field
-            cursor.execute('CREATE INDEX IF NOT EXISTS idx_access_key ON emails (access_key)')
-            
             conn.commit()
         except Exception as e:
             conn.rollback()
@@ -83,17 +81,8 @@ class EmailDatabase:
         """Generate a unique ID for the email record."""
         return str(uuid.uuid4())
     
-    def _generate_access_key(self, email_id: str) -> str:
-        """
-        Generate an access key for retrieving the original email.
-        This acts as a security measure to prevent unauthorized access.
-        """
-        # Use a combination of the email ID and a timestamp, hashed
-        data = f"{email_id}:{datetime.now().isoformat()}:{os.urandom(8).hex()}"
-        return hashlib.sha256(data.encode()).hexdigest()
-    
     def store_email(self, original_email: str, masked_email: str, 
-                   masked_entities: List[Dict[str, Any]], category: Optional[str] = None) -> Tuple[str, str]:
+                   masked_entities: List[Dict[str, Any]], category: Optional[str] = None) -> str:
         """
         Store the original email along with its masked version and related information.
         
@@ -104,32 +93,30 @@ class EmailDatabase:
             category: Optional category of the email
             
         Returns:
-            Tuple of (email_id, access_key) for future reference
+            email_id for future reference
         """
         conn = self._get_connection()
         try:
             cursor = conn.cursor()
             
             email_id = self._generate_id()
-            access_key = self._generate_access_key(email_id)
             
             # Store the email data
             cursor.execute(
-                'INSERT INTO emails (id, original_email, masked_email, masked_entities, category, created_at, access_key) '
-                'VALUES (?, ?, ?, ?, ?, ?, ?)',
+                'INSERT INTO emails (id, original_email, masked_email, masked_entities, category, created_at) '
+                'VALUES (?, ?, ?, ?, ?, ?)',
                 (
                     email_id,
                     original_email,
                     masked_email,
                     json.dumps(masked_entities),  # Convert to JSON string for SQLite
                     category,
-                    datetime.now().isoformat(),
-                    access_key
+                    datetime.now().isoformat()
                 )
             )
             
             conn.commit()
-            return email_id, access_key
+            return email_id
         except Exception as e:
             conn.rollback()
             raise e
@@ -147,14 +134,18 @@ class EmailDatabase:
         Returns:
             Dictionary with email data or None if not found or access_key is invalid
         """
+        # Verify the access key matches the global access key
+        if access_key != self.access_key:
+            return None
+            
         conn = self._get_connection()
         try:
             cursor = conn.cursor()
             
             cursor.execute(
                 'SELECT id, original_email, masked_email, masked_entities, category, created_at '
-                'FROM emails WHERE id = ? AND access_key = ?',
-                (email_id, access_key)
+                'FROM emails WHERE id = ?',
+                (email_id,)
             )
             
             row = cursor.fetchone()
@@ -203,5 +194,40 @@ class EmailDatabase:
                 "category": row[3],
                 "created_at": row[4]
             }
+        finally:
+            conn.close()
+    
+    def get_email_by_masked_content(self, masked_email: str) -> Optional[Dict[str, Any]]:
+        """
+        Retrieve the original email using the masked email content.
+        
+        Args:
+            masked_email: The masked version of the email to search for
+            
+        Returns:
+            Dictionary with full email data or None if not found
+        """
+        conn = self._get_connection()
+        try:
+            cursor = conn.cursor()
+            
+            cursor.execute(
+                'SELECT id, original_email, masked_email, masked_entities, category, created_at '
+                'FROM emails WHERE masked_email = ?',
+                (masked_email,)
+            )
+            
+            row = cursor.fetchone()
+            if not row:
+                return None
+                
+            return {
+                "id": row[0],
+                "original_email": row[1],
+                "masked_email": row[2],
+                "masked_entities": json.loads(row[3]),  # Convert from JSON string back to Python dict
+                "category": row[4],
+                "created_at": row[5]
+            }
         finally:
             conn.close()

main.py

@@ -46,9 +46,9 @@ class EmailOutput(BaseModel):
     masked_email: str
     category_of_the_email: str
 
-class EmailRetrievalInput(BaseModel):
-    """Input model for retrieving original email"""
-    email_id: str
+class MaskedEmailInput(BaseModel):
+    """Input model for retrieving original email by masked email content"""
+    masked_email: str
     access_key: str
 
 @app.post("/classify", response_model=EmailOutput)
@@ -63,7 +63,7 @@ async def classify_email(email_input: EmailInput) -> Dict[str, Any]:
         The classified email data with masked PII
     """
     try:
-        # Process the email to mask PII
+        # Process the email to mask PII and store original in database
         processed_data = pii_masker.process_email(email_input.input_email_body)
         
         # Classify the masked email
@@ -79,35 +79,44 @@ async def classify_email(email_input: EmailInput) -> Dict[str, Any]:
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error processing email: {str(e)}")
 
-@app.post("/api/v1/original-email/retrieve", response_model=Dict[str, Any])
-async def retrieve_original_email_v1(retrieval_input: EmailRetrievalInput) -> Dict[str, Any]:
+@app.post("/api/v1/unmask-email", response_model=Dict[str, Any])
+async def unmask_email(masked_email_input: MaskedEmailInput) -> Dict[str, Any]:
     """
-    New API endpoint to retrieve the original unmasked email from SQLite database.
+    Retrieve the original unmasked email using the masked email content from the classify response.
     
     Args:
-        retrieval_input: The email ID and access key
+        masked_email_input: Contains the masked email and access key
         
     Returns:
         The original email data with PII information
     """
     try:
-        email_data = pii_masker.get_original_email(
-            retrieval_input.email_id, 
-            retrieval_input.access_key
-        )
+        # Verify access key matches the global access key
+        if masked_email_input.access_key != os.environ.get("EMAIL_ACCESS_KEY", "default_secure_access_key"):
+            raise HTTPException(status_code=401, detail="Invalid access key")
+        
+        # Retrieve the original email using the masked content
+        email_data = pii_masker.get_original_by_masked_email(masked_email_input.masked_email)
         
         if not email_data:
-            raise HTTPException(status_code=404, detail="Email not found or invalid access key")
+            raise HTTPException(status_code=404, detail="Original email not found for the provided masked email")
         
         return {
             "status": "success",
-            "data": email_data,
+            "data": {
+                "id": email_data["id"],
+                "original_email": email_data["original_email"],
+                "masked_email": email_data["masked_email"],
+                "masked_entities": email_data["masked_entities"],
+                "category": email_data.get("category", ""),
+                "created_at": email_data.get("created_at", "")
+            },
             "message": "Original email retrieved successfully"
         }
     except Exception as e:
         if isinstance(e, HTTPException):
             raise e
-        raise HTTPException(status_code=500, detail=f"Error retrieving email: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error retrieving original email: {str(e)}")
 
 @app.get("/health")
 async def health_check():

utils.py

@@ -329,21 +329,20 @@ class PIIMasker:
         # Mask the email
         masked_email, entity_info = self.mask_text(email_text)
         
-        # Store the email in the SQLite database
-        email_id, access_key = self.db.store_email(
+        # Store the email in the SQLite database - only get back email_id now
+        email_id = self.db.store_email(
             original_email=email_text,
             masked_email=masked_email,
             masked_entities=entity_info
         )
         
-        # Return the processed data with database references
+        # Return the processed data with just the email_id
         return {
-            "input_email_body": email_text,  # Return original input for compatibility
+            "input_email_body": email_text,  # Return original input for API compatibility
             "list_of_masked_entities": entity_info,
             "masked_email": masked_email,
             "category_of_the_email": "",
-            "email_id": email_id,
-            "access_key": access_key  # Include access key for immediate retrieval if needed
+            "email_id": email_id
         }
     
     def get_original_email(self, email_id: str, access_key: str) -> Optional[Dict[str, Any]]:
@@ -369,4 +368,16 @@ class PIIMasker:
         Returns:
             The masked email data or None if not found
         """
-        return self.db.get_email_by_id(email_id)
+        return self.db.get_email_by_id(email_id)
+        
+    def get_original_by_masked_email(self, masked_email: str) -> Optional[Dict[str, Any]]:
+        """
+        Retrieve the original unmasked email using the masked email content.
+        
+        Args:
+            masked_email: The masked version of the email to search for
+            
+        Returns:
+            The original email data or None if not found
+        """
+        return self.db.get_email_by_masked_content(masked_email)