"""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()
