SQLValidator.py

import re

class SQLQueryValidator():
    def __init__(self):
        pass
        
    # list of DDL DML etc and other keywords that can be used for abusing db
    prohibited_sql_keywords = [
        'ALTER', 'RENAME', 'DELETE', 'UPDATE', 'CREATE', 'TRUNCATE', 'RENAME', 'GRANT', 'COMMENT', 'REVOKE', 'ANALYZE',
        'REINDEX', 'CLUSTER', 'INSERT', 'MERGE', 'REPLACE', 'COMMIT', 'ROLLBACK', 'SAVEPOINT', 'SET TRANSACTION', 'SET',
        'RESET', 'DISCARD', 'EXEC', 'EXECUTE', 'CALL', 'SHOW', 'LOAD', 'COPY', '1=1', 'pg_sleep', 'BENCHMARK', 
        'LOAD_FILE', 'SLEEP'
    ]
    
    def _is_query_safe(self, query : str) -> dict[str, str, bool] :
        '''
        checks if the query is safe or not
        '''
        if len(query) > 1000:
            return {
                "status" : "DANGEROUS",
                "reason" : f"query length is greater than 1000 characters ({len(query)})",
                "safe" : False
            }
        
        query = query.upper()
        
        if r'%25' in query or r'%5C' in query:
            return {
                "status" : "DANGEROUS",
                "reason" : "query contains escape characters",
                "safe" : False
            }
        
        for k in self.prohibited_sql_keywords:
            if re.search(rf'\b{k}\b', query):
                return {
                    "status" : "DANGEROUS",
                    "reason" : f"query contains prohibited keyword {k}",
                    "safe" : False
                }
        return {
            "status" : "SAFE",
            "reason" : "no potentially dangerous keywords found",
            "safe" : True
        }
    
    def check_query(self, sql_query : str) -> bool :
        res = self._is_query_safe(sql_query)
        return True if res["safe"] else False