Initial commit.

mail/spamassassin/pyzor-0.7.0/tests/unit/test_server.py

@@ -0,0 +1,341 @@
+"""Test the pyzor.server module
+"""
+import io
+import sys
+import time
+import logging
+import hashlib
+import unittest
+import SocketServer
+
+from datetime import datetime, timedelta
+
+import pyzor
+import pyzor.server
+import pyzor.engines.common
+
+class MockServer():
+    """Mocks the pyzor.server.Server class"""
+    
+    def __init__(self):
+        self.log = logging.getLogger("pyzord")
+        self.usage_log = logging.getLogger("pyzord-usage")
+        self.log.addHandler(logging.NullHandler())
+        self.usage_log.addHandler(logging.NullHandler())
+        
+        
+class MockDatagramRequestHandler():
+    """ Mock the SocketServer.DatagramRequestHand."""
+    
+    def __init__(self, headers, database=None, acl=None, accounts=None):
+        """Initiates an request handler and set's the data in `headers` as 
+        the request. Also set's the database, acl and accounts for the 
+        MockServer. 
+        
+        This will be set as base class for RequestHandler. 
+        """
+        self.rfile = io.BytesIO()
+        self.wfile = io.BytesIO()
+        for i, j in headers.iteritems():
+            self.rfile.write(("%s: %s\n" % (i, j)).encode("utf8"))
+        self.rfile.seek(0)
+        self.packet = None
+        self.client_address = ["127.0.0.1"]        
+        
+        # Setup MockServer data
+        self.server = MockServer()
+        self.server.database = database
+        if acl:
+            self.server.acl = acl
+        else:
+            self.server.acl = {pyzor.anonymous_user: ("check", "report", "ping", "info", "whitelist",)}
+        self.server.accounts = accounts                
+        
+        self.handle()
+            
+    def handle(self):
+        pass
+
+
+class RequestHandlerTest(unittest.TestCase):            
+    
+    def setUp(self):
+        unittest.TestCase.setUp(self)
+        self.real_drh = SocketServer.DatagramRequestHandler
+        SocketServer.DatagramRequestHandler = MockDatagramRequestHandler
+        pyzor.server.RequestHandler.__bases__ = (MockDatagramRequestHandler,)
+        
+        # setup the basic values for request and response
+        self.request = {"User": pyzor.anonymous_user,
+                        "Time": str(int(time.time())),
+                        "PV": str(pyzor.proto_version),
+                        "Thread": "3597"}
+        self.expected_response = {"Code": "200",
+                                  "Diag": "OK",
+                                  "PV": str(pyzor.proto_version),
+                                  "Thread": "3597"}
+        
+    def tearDown(self):
+        unittest.TestCase.tearDown(self)
+        SocketServer.DatagramRequestHandler = self.real_drh
+        pyzor.server.RequestHandler.__bases__ = (self.real_drh,)
+    
+    def check_response(self, handler):
+        """Checks if the response from the handler is equal to
+        the expected response.
+        """
+        handler.wfile.seek(0)
+        response = handler.wfile.read()
+        response = response.decode("utf8").replace("\n\n", "\n")
+        
+        result = {}
+        for line in response.splitlines():
+            key = line.split(":", 1)[0].strip()
+            value = line.split(":")[1].strip()
+            result[key] = value
+        
+        self.assertEqual(result, self.expected_response)            
+    
+    def timestamp(self, time_obj):
+        if not time_obj:
+            return 0
+        else:
+            return str(int(time.mktime(time_obj.timetuple())))
+    
+    def test_ping(self):
+        """Tests the ping command handler"""
+        self.request["Op"] = "ping"
+        handler = pyzor.server.RequestHandler(self.request)
+        
+        self.check_response(handler)
+    
+    def test_pong(self):
+        """Tests the pong command handler"""
+        digest = "2aedaac999d71421c9ee49b9d81f627a7bc570aa"
+        database = {digest: pyzor.engines.common.Record(24, 42)}
+        
+        self.request["Op"] = "pong"
+        self.request["Op-Digest"] = digest
+        handler = pyzor.server.RequestHandler(self.request, database)
+        self.expected_response["Count"] = str(sys.maxint) 
+        self.expected_response["WL-Count"] = "0"
+
+    def test_check(self):
+        """Tests the check command handler"""
+        digest = "2aedaac999d71421c9ee49b9d81f627a7bc570aa"
+        database = {digest: pyzor.engines.common.Record(24, 42)}
+        
+        self.request["Op"] = "check"
+        self.request["Op-Digest"] = digest
+        handler = pyzor.server.RequestHandler(self.request, database)
+        self.expected_response["Count"] = "24"
+        self.expected_response["WL-Count"] = "42"
+        
+        self.check_response(handler)
+        
+    def test_check_new(self):
+        """Tests the check command handler with a new record"""
+        digest = "2aedaac999d71421c9ee49b9d81f627a7bc570aa"
+        database = {}
+        
+        self.request["Op"] = "check"
+        self.request["Op-Digest"] = digest
+        handler = pyzor.server.RequestHandler(self.request, database)
+        self.expected_response["Count"] = "0"
+        self.expected_response["WL-Count"] = "0"
+        
+        self.check_response(handler)
+    
+    def test_info(self):
+        """Tests the info command handler"""
+        entered = datetime.now() - timedelta(days=10)
+        updated = datetime.now()
+        wl_entered = datetime.now() - timedelta(days=20)
+        wl_updated = datetime.now() - timedelta(days=2)
+                
+        digest = "2aedaac999d71421c9ee49b9d81f627a7bc570aa"
+        database = {digest: pyzor.engines.common.Record(24, 42, entered, updated,
+                                                        wl_entered, wl_updated)}        
+        self.request["Op"] = "info"
+        self.request["Op-Digest"] = digest
+        handler = pyzor.server.RequestHandler(self.request, database)
+        self.expected_response["Count"] = "24"
+        self.expected_response["WL-Count"] = "42"
+        self.expected_response["Entered"] = self.timestamp(entered)
+        self.expected_response["Updated"] = self.timestamp(updated)
+        self.expected_response["WL-Entered"] = self.timestamp(wl_entered)
+        self.expected_response["WL-Updated"] = self.timestamp(wl_updated)
+        
+        self.check_response(handler)
+        
+    def test_info_new(self):
+        """Tests the info command handler with a new record"""
+        digest = "2aedaac999d71421c9ee49b9d81f627a7bc570aa"
+        database = {}        
+        self.request["Op"] = "info"
+        self.request["Op-Digest"] = digest
+        handler = pyzor.server.RequestHandler(self.request, database)
+        self.expected_response["Count"] = "0"
+        self.expected_response["WL-Count"] = "0"
+        self.expected_response["Entered"] = "0"
+        self.expected_response["Updated"] = "0"
+        self.expected_response["WL-Entered"] = "0"
+        self.expected_response["WL-Updated"] = "0"
+        
+        self.check_response(handler)
+        
+    def test_report(self):
+        """Tests the report command handler"""
+        digest = "2aedaac999d71421c9ee49b9d81f627a7bc570aa"
+        database = {digest: pyzor.engines.common.Record(24, 42)}
+         
+        self.request["Op"] = "report"
+        self.request["Op-Digest"] = digest
+        handler = pyzor.server.RequestHandler(self.request, database)
+        
+        self.check_response(handler)
+        self.assertEqual(database[digest].r_count, 25)            
+        
+    def test_report_new(self):
+        """Tests the report command handler with a new record"""
+        digest = "2aedaac999d71421c9ee49b9d81f627a7bc570aa"
+        database = {}
+         
+        self.request["Op"] = "report"
+        self.request["Op-Digest"] = digest
+        handler = pyzor.server.RequestHandler(self.request, database)
+        
+        self.check_response(handler)
+        self.assertEqual(database[digest].r_count, 1)   
+        
+    def test_whitelist(self):
+        """Tests the whitelist command handler"""
+        digest = "2aedaac999d71421c9ee49b9d81f627a7bc570aa"
+        database = {digest: pyzor.engines.common.Record(24, 42)}
+         
+        self.request["Op"] = "whitelist"
+        self.request["Op-Digest"] = digest
+        handler = pyzor.server.RequestHandler(self.request, database)
+        
+        self.check_response(handler)
+        self.assertEqual(database[digest].wl_count, 43)            
+        
+    def test_whitelist_new(self):
+        """Tests the whitelist command handler with a new record"""
+        digest = "2aedaac999d71421c9ee49b9d81f627a7bc570aa"
+        database = {}
+         
+        self.request["Op"] = "whitelist"
+        self.request["Op-Digest"] = digest
+        handler = pyzor.server.RequestHandler(self.request, database)
+        
+        self.check_response(handler)
+        self.assertEqual(database[digest].wl_count, 1)
+    
+    def test_handle_no_version(self):
+        """Tests handling an request with no version specified"""
+        self.request["Op"] = "ping"
+        del self.request["PV"]
+        handler = pyzor.server.RequestHandler(self.request)
+        
+        self.expected_response["Code"] = "400"
+        self.expected_response["Diag"] = "Bad request"
+        self.check_response(handler)
+    
+    def test_handle_unsupported_version(self):
+        """Tests handling an request with an unsupported version specified"""
+        self.request["Op"] = "ping"
+        self.request["PV"] = str(pyzor.proto_version + 2)
+        handler = pyzor.server.RequestHandler(self.request)
+        
+        self.expected_response["Code"] = "505"
+        self.expected_response["Diag"] = "Version Not Supported"
+        self.check_response(handler)        
+    
+    def test_handle_not_implemented(self):
+        """Tests handling an request with an unimplemented command"""
+        self.request["Op"] = "notimplemented"
+        acl = {pyzor.anonymous_user: "notimplemented"}
+        handler = pyzor.server.RequestHandler(self.request, acl=acl)
+        
+        self.expected_response["Code"] = "501"
+        self.expected_response["Diag"] = "Not implemented"
+        self.check_response(handler)
+
+    def test_handle_unauthorized(self):
+        """Tests handling an request with an unauthorized command"""
+        self.request["Op"] = "report"
+        acl = {pyzor.anonymous_user: ("ping", "check")}
+        handler = pyzor.server.RequestHandler(self.request, acl=acl)
+        
+        self.expected_response["Code"] = "403"
+        self.expected_response["Diag"] = "Forbidden"
+        self.check_response(handler)
+        
+    def test_handle_account(self):
+        """Tests handling an request where user is not anonymous"""
+        self.request["Op"] = "ping"
+        self.request["User"] = "testuser"
+        acl = {"testuser": ("ping", "check")}
+        accounts = {"testuser": "testkey"}
+                
+        mock_vs = lambda x, y: None
+        real_vs = pyzor.account.verify_signature
+        pyzor.account.verify_signature = mock_vs
+        try:
+            handler = pyzor.server.RequestHandler(self.request, acl=acl,
+                                                  accounts=accounts)
+            self.check_response(handler)
+        finally:
+            pyzor.account.verify_signature = real_vs
+    
+    def test_handle_unknown_account(self):
+        """Tests handling an request where user is unkwown"""
+        self.request["Op"] = "ping"
+        self.request["User"] = "testuser"
+        acl = {"testuser": ("ping", "check")}
+        accounts = {}
+        
+        self.expected_response["Code"] = "401"
+        self.expected_response["Diag"] = "Unauthorized"
+        
+        def mock_vs(x, y):
+            pass        
+        real_vs = pyzor.account.verify_signature
+        pyzor.account.verify_signature = mock_vs
+        try:
+            handler = pyzor.server.RequestHandler(self.request, acl=acl,
+                                                  accounts=accounts)
+            self.check_response(handler)
+        finally:
+            pyzor.account.verify_signature = real_vs
+
+    def test_handle_invalid_signature(self):
+        """Tests handling an request where user key is invalid"""
+        self.request["Op"] = "ping"
+        self.request["User"] = "testuser"
+        acl = {"testuser": ("ping", "check")}
+        accounts = {"testuser": ("ping", "check")}
+        
+        self.expected_response["Code"] = "401"
+        self.expected_response["Diag"] = "Unauthorized"
+        
+        def mock_vs(x, y):
+            raise pyzor.SignatureError("Invalid signature.")        
+        real_vs = pyzor.account.verify_signature
+        pyzor.account.verify_signature = mock_vs
+        try:
+            handler = pyzor.server.RequestHandler(self.request, acl=acl,
+                                                  accounts=accounts)
+            self.check_response(handler)
+        finally:
+            pyzor.account.verify_signature = real_vs
+
+def suite():
+    """Gather all the tests from this module in a test suite."""
+    test_suite = unittest.TestSuite()
+    test_suite.addTest(unittest.makeSuite(RequestHandlerTest))
+    return test_suite
+
+if __name__ == '__main__':
+    unittest.main()