# Path: protector/protector_main.py
# class Protector(object):
#     """
#     The main protector class which checks for malicious queries
#     """
# 
#     def __init__(self, rules, whitelist=[], safe_mode=True):
#         """
#         :param rules: A list of rules to evaluate
#         :param safe_mode: If set to True, allow the query in case it can not be parsed
#         :return:
#         """
#         self.parser = QueryParser()
#         self.guard = Guard(rules)
#         self.sanitizer = Sanitizer()
#         self.whitelist = whitelist
#         self.safe_mode = safe_mode
# 
#     def check(self, query_string):
#         logging.debug("Checking {}".format(query_string))
#         query_sanitized = self.sanitizer.sanitize(query_string)
#         query = self.parser.parse(query_sanitized)
#         if query:
#             if self.is_whitelisted(query):
#                 return Ok(True)
#             return self.guard.is_allowed(query)
#         else:
#             error_msg = "Could not parse query: '{}'".format(query_string)
#             logging.info(error_msg)
#             if self.safe_mode:
#                 return Ok(True)
#             else:
#                 return Err(error_msg)
# 
#     def is_whitelisted(self, query):
#         if query.get_type() in {Keyword.LIST, Keyword.DROP}:
#             series = query.series_stmt
#         else:
#             series = query.from_stmt
#         for pattern in self.whitelist:
#             match = re.match(pattern, series)
#             if match:
#                 return True
#         return False


import unittest
from urllib2 import quote
from protector.protector_main import Protector



class TestProtector(unittest.TestCase):
    def test_whitelist(self):
        protector = Protector(["too_many_datapoints"], ["^releases$", "^grafana\.", ".*java.*boot.*version.*"], True)
        self.assertTrue(protector.check(quote("select * from releases")).is_ok())
        self.assertFalse(protector.check(quote("select * from some.random.series.releases")).is_ok())
        self.assertFalse(protector.check(quote("select * from releases.are.important")).is_ok())

        self.assertTrue(protector.check(quote("select * from grafana.temp_dashboard_Y2hyb25vcy1qb2Jz")).is_ok())
        self.assertFalse(protector.check(quote("select * from bla.bla.grafana.")).is_ok())
        self.assertFalse(protector.check(quote("select * from grafana")).is_ok())

        query = "select tags,title,text_details from merge(/app4\.java\.boot\.version\./) where time > now() - 7d"
        self.assertTrue(protector.check(quote(query)).is_ok())

        protector = Protector(["too_many_datapoints"], [], True)
        self.assertFalse(protector.check(quote("select * from grafana")).is_ok())

    def test_unknown_queries_safe_mode(self):
        protector = Protector(["prevent_delete"], [], True)
        self.assertTrue(protector.check("").is_ok())
        self.assertTrue(protector.check("asdf").is_ok())

    def test_unknown_queries_non_safe_mode(self):

# Path: protector/protector_main.py
# class Protector(object):
#     """
#     The main protector class which checks for malicious queries
#     """
# 
#     def __init__(self, rules, whitelist=[], safe_mode=True):
#         """
#         :param rules: A list of rules to evaluate
#         :param safe_mode: If set to True, allow the query in case it can not be parsed
#         :return:
#         """
#         self.parser = QueryParser()
#         self.guard = Guard(rules)
#         self.sanitizer = Sanitizer()
#         self.whitelist = whitelist
#         self.safe_mode = safe_mode
# 
#     def check(self, query_string):
#         logging.debug("Checking {}".format(query_string))
#         query_sanitized = self.sanitizer.sanitize(query_string)
#         query = self.parser.parse(query_sanitized)
#         if query:
#             if self.is_whitelisted(query):
#                 return Ok(True)
#             return self.guard.is_allowed(query)
#         else:
#             error_msg = "Could not parse query: '{}'".format(query_string)
#             logging.info(error_msg)
#             if self.safe_mode:
#                 return Ok(True)
#             else:
#                 return Err(error_msg)
# 
#     def is_whitelisted(self, query):
#         if query.get_type() in {Keyword.LIST, Keyword.DROP}:
#             series = query.series_stmt
#         else:
#             series = query.from_stmt
#         for pattern in self.whitelist:
#             match = re.match(pattern, series)
#             if match:
#                 return True
#         return False


# Path: protector/tests/protector_test/test_protector.py
import unittest
from urllib2 import quote
from protector.protector_main import Protector




class TestProtector(unittest.TestCase):
    def test_whitelist(self):
        protector = Protector(["too_many_datapoints"], ["^releases$", "^grafana\.", ".*java.*boot.*version.*"], True)
        self.assertTrue(protector.check(quote("select * from releases")).is_ok())
        self.assertFalse(protector.check(quote("select * from some.random.series.releases")).is_ok())
        self.assertFalse(protector.check(quote("select * from releases.are.important")).is_ok())

        self.assertTrue(protector.check(quote("select * from grafana.temp_dashboard_Y2hyb25vcy1qb2Jz")).is_ok())
        self.assertFalse(protector.check(quote("select * from bla.bla.grafana.")).is_ok())
        self.assertFalse(protector.check(quote("select * from grafana")).is_ok())

        query = "select tags,title,text_details from merge(/app4\.java\.boot\.version\./) where time > now() - 7d"
        self.assertTrue(protector.check(quote(query)).is_ok())

        protector = Protector(["too_many_datapoints"], [], True)
        self.assertFalse(protector.check(quote("select * from grafana")).is_ok())

    def test_unknown_queries_safe_mode(self):
        protector = Protector(["prevent_delete"], [], True)
        self.assertTrue(protector.check("").is_ok())
        self.assertTrue(protector.check("asdf").is_ok())

    def test_unknown_queries_non_safe_mode(self):

        protector = Protector(["prevent_delete"], [], False)