import unittest
from StringIO import StringIO
from mlspolicy import mlspolicy, domrange, ClearanceNode, ResourceEntry, UserEntry

source = '''
clearances: clr1 > clr2 > clr4
# whole line comment
clearances: clr2 > clr5         # foo
clearances: clr1 > clr3 > clr5
categories: cat1, cat2            # bar buzz 

assign: clr1:cat1 -r /usr        # a comment
assign: clr3:cat2 /usr/somefile

users: clr2:cat1,cat2 user1,user2
users: clr3:cat1 user3,user4
'''
clearance_lattice_ref = {'clr1': ClearanceNode(set(), {'clr2', 'clr3'}),
                         'clr2': ClearanceNode({'clr1'}, {'clr4', 'clr5'}),
                         'clr3': ClearanceNode({'clr1'}, {'clr5'}),
                         'clr4': ClearanceNode({'clr2'}, set()),
                         'clr5': ClearanceNode({'clr2', 'clr3'}, set())}

categories_ref = {'cat1', 'cat2'}

resource_dict_ref = {'/usr': {ResourceEntry('clr1', 'cat1', True)},
                     '/usr/somefile': {ResourceEntry('clr3', 'cat2', False)}}

user_dict_ref = {'user1': {UserEntry('clr2', 'cat1'), UserEntry('clr2', 'cat2')},
                 'user2': {UserEntry('clr2', 'cat1'), UserEntry('clr2', 'cat2')},
                 'user3': {UserEntry('clr3', 'cat1')},
                 'user4': {UserEntry('clr3', 'cat1')}}

class TestMLS(unittest.TestCase):
    def setUp(self):
        (self.clearance_lattice,
         self.categories,
         self.resource_dict,
         self.user_dict) = mlspolicy(StringIO(source))

    def test_lattice_eq(self):
        self.assertEqual(clearance_lattice_ref, self.clearance_lattice)

    def test_categories_eq(self):
        self.assertEqual(categories_ref, self.categories)
    
    def test_resource_dict_eq(self):
        self.assertEqual(resource_dict_ref, self.resource_dict)
    
    def test_user_dict_eq(self):
        self.assertEqual(user_dict_ref, self.user_dict)

    def test_domrange(self):
        # clr5 on bottom
        result = domrange('clr5', 'clr5', self.clearance_lattice)
        self.assertEqual({'clr5'}, result)

        result = domrange('clr4', 'clr5', self.clearance_lattice)
        self.assertEqual(set(), result)
        
        result = domrange('clr3', 'clr5', self.clearance_lattice)
        self.assertEqual({'clr3', 'clr5'}, result)
        
        result = domrange('clr2', 'clr5', self.clearance_lattice)
        self.assertEqual({'clr2', 'clr5'}, result)
        
        result = domrange('clr1', 'clr5', self.clearance_lattice)
        self.assertEqual({'clr5', 'clr3', 'clr2', 'clr1'}, result)
        
        # clr4 on bottom
        result = domrange('clr5', 'clr4', self.clearance_lattice)
        self.assertEqual(set(), result)

        result = domrange('clr4', 'clr4', self.clearance_lattice)
        self.assertEqual({'clr4'}, result)
        
        result = domrange('clr3', 'clr4', self.clearance_lattice)
        self.assertEqual(set(), result)
        
        result = domrange('clr2', 'clr4', self.clearance_lattice)
        self.assertEqual({'clr2', 'clr4'}, result)
        
        result = domrange('clr1', 'clr4', self.clearance_lattice)
        self.assertEqual({'clr4', 'clr2', 'clr1'}, result)
        
        # clr3 on bottom
        result = domrange('clr5', 'clr3', self.clearance_lattice)
        self.assertEqual(set(), result)

        result = domrange('clr4', 'clr3', self.clearance_lattice)
        self.assertEqual(set(), result)
        
        result = domrange('clr3', 'clr3', self.clearance_lattice)
        self.assertEqual({'clr3'}, result)
        
        result = domrange('clr2', 'clr3', self.clearance_lattice)
        self.assertEqual(set(), result)
        
        result = domrange('clr1', 'clr3', self.clearance_lattice)
        self.assertEqual({'clr3', 'clr1'}, result)
        
        # clr2 on bottom
        result = domrange('clr5', 'clr2', self.clearance_lattice)
        self.assertEqual(set(), result)

        result = domrange('clr4', 'clr2', self.clearance_lattice)
        self.assertEqual(set(), result)
        
        result = domrange('clr3', 'clr2', self.clearance_lattice)
        self.assertEqual(set(), result)
        
        result = domrange('clr2', 'clr2', self.clearance_lattice)
        self.assertEqual({'clr2'}, result)
        
        result = domrange('clr1', 'clr2', self.clearance_lattice)
        self.assertEqual({'clr2', 'clr1'}, result)

        # clr1 on bottom
        result = domrange('clr5', 'clr1', self.clearance_lattice)
        self.assertEqual(set(), result)

        result = domrange('clr4', 'clr1', self.clearance_lattice)
        self.assertEqual(set(), result)
        
        result = domrange('clr3', 'clr1', self.clearance_lattice)
        self.assertEqual(set(), result)
        
        result = domrange('clr2', 'clr1', self.clearance_lattice)
        self.assertEqual(set(), result)
        
        result = domrange('clr1', 'clr1', self.clearance_lattice)
        self.assertEqual({'clr1'}, result)


if __name__ == '__main__':
    unittest.main()
