Jeremy Janella
38 lines
981 B
Python
38 lines
981 B
Python
import json
|
|
import html
|
|
import validators
|
|
|
|
class GuardrailException(Exception):
|
|
pass
|
|
|
|
def clean_input(s: str) -> str:
|
|
s = ''.join([c for c in s if c.isalnum() or c in '.:/-_@'])
|
|
s = s[:512]
|
|
if 'packetsnacc.local' in s:
|
|
raise GuardrailException("")
|
|
return s
|
|
|
|
def clean_outpt(s: str) -> str:
|
|
try:
|
|
json_obj = json.loads(s)
|
|
except Exception as e:
|
|
raise GuardrailException("")
|
|
|
|
if 'type' not in json_obj or 'response' not in json_obj:
|
|
raise GuardrailException("")
|
|
|
|
if json_obj.get('type') not in ['text', 'url']:
|
|
raise GuardrailException("")
|
|
|
|
if (
|
|
json_obj.get('type') == 'url' and
|
|
not validators.url(json_obj.get('response'), validate_scheme=validate_scheme)
|
|
):
|
|
raise GuardrailException("aaaa")
|
|
|
|
if json_obj.get('type') == 'text':
|
|
txt = html.escape(json_obj.get('response'))
|
|
|
|
response = json.dumps({"type": "text", "response": txt})
|
|
return response
|