prelington commited on
Commit
c03f547
·
verified ·
1 Parent(s): f592f9a

Create security.py

Browse files
Files changed (1) hide show
  1. security.py +39 -0
security.py ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import re
2
+ from typing import List
3
+
4
+ class SecurityValidator:
5
+ @staticmethod
6
+ def validate_input(text: str, max_length: int = 1000) -> bool:
7
+ """Validate input text for security"""
8
+ if len(text) > max_length:
9
+ raise ValueError(f"Input too long. Max {max_length} characters.")
10
+
11
+ # Check for potential prompt injection
12
+ injection_patterns = [
13
+ r"ignore previous",
14
+ r"system prompt",
15
+ r"forget your instructions"
16
+ ]
17
+
18
+ for pattern in injection_patterns:
19
+ if re.search(pattern, text, re.IGNORECASE):
20
+ raise SecurityError("Potential prompt injection detected")
21
+
22
+ return True
23
+
24
+ @staticmethod
25
+ def sanitize_output(text: str) -> str:
26
+ """Sanitize model output"""
27
+ # Remove any potentially harmful content
28
+ harmful_patterns = [
29
+ r"<script.*?>.*?</script>",
30
+ r"<iframe.*?>.*?</iframe>"
31
+ ]
32
+
33
+ for pattern in harmful_patterns:
34
+ text = re.sub(pattern, "", text, flags=re.IGNORECASE | re.DOTALL)
35
+
36
+ return text
37
+
38
+ class SecurityError(Exception):
39
+ pass