#!/usr/bin/env python # -*- coding: utf-8 -*- # Originally based on ./sam.py import optparse import sys import os import base64 import random import re sys.path.insert(0, "bin/python") import samba from samba.tests.subunitrun import SubunitOptions, TestProgram import samba.getopt as options from samba.auth import system_session import ldb from samba.samdb import SamDB import time parser = optparse.OptionParser("vlv.py [options] ") sambaopts = options.SambaOptions(parser) parser.add_option_group(sambaopts) parser.add_option_group(options.VersionOptions(parser)) # use command line creds if available credopts = options.CredentialsOptions(parser) parser.add_option_group(credopts) subunitopts = SubunitOptions(parser) parser.add_option_group(subunitopts) parser.add_option('--elements', type='int', default=20, help="use this many elements in the tests") parser.add_option('--delete-in-setup', action='store_true', help="cleanup in next setup rather than teardown") parser.add_option('--skip-attr-regex', help="ignore attributes matching this regex") opts, args = parser.parse_args() if len(args) < 1: parser.print_usage() sys.exit(1) host = args[0] lp = sambaopts.get_loadparm() creds = credopts.get_credentials(lp) N_ELEMENTS = opts.elements class VlvTestException(Exception): pass def encode_vlv_control(critical=1, before=0, after=0, offset=None, gte=None, n=0, cookie=None): s = "vlv:%d:%d:%d:" % (critical, before, after) if offset is not None: m = "%d:%d" % (offset, n) elif ':' in gte or '\x00' in gte: gte = base64.b64encode(gte) m = "base64>=%s" % gte else: m = ">=%s" % gte if cookie is None: return s + m return s + m + ':' + cookie def get_cookie(controls, expected_n=None): """Get the cookie, STILL base64 encoded, or raise ValueError.""" for c in list(controls): cstr = str(c) if cstr.startswith('vlv_resp'): head, n, _, cookie = cstr.rsplit(':', 3) if expected_n is not None and int(n) != expected_n: raise ValueError("Expected %s items, server said %s" % (expected_n, n)) return cookie raise ValueError("there is no VLV response") class VLVTests(samba.tests.TestCase): def create_user(self, i, n, prefix='vlvtest', suffix='', attrs=None): name = "%s%d%s" % (prefix, i, suffix) user = { 'cn': name, "objectclass": "user", 'givenName': "abcdefghijklmnopqrstuvwxyz"[i % 26], "roomNumber": "%sbc" % (n - i), "carLicense": "后来经", "employeeNumber": "%s%sx" % (abs(i * (99 - i)), '\n' * (i & 255)), "accountExpires": "%s" % (10 ** 9 + 1000000 * i), "msTSExpireDate4": "19%02d0101010000.0Z" % (i % 100), "flags": str(i * (n - i)), "serialNumber": "abc %s%s%s" % ('AaBb |-/'[i & 7], ' 3z}'[i & 3], '"@'[i & 1],), } # _user_broken_attrs tests are broken due to problems outside # of VLV. _user_broken_attrs = { # Sort doesn't look past a NUL byte. "photo": "\x00%d" % (n - i), "audio": "%sn octet string %s%s ♫♬\x00lalala" % ('Aa'[i & 1], chr(i & 255), i), "displayNamePrintable": "%d\x00%c" % (i, i & 255), "adminDisplayName": "%d\x00b" % (n-i), "title": "%d%sb" % (n - i, '\x00' * i), "comment": "Favourite colour is %d" % (n % (i + 1)), # Names that vary only in case. Windows returns # equivalent addresses in the order they were put # in ('a st', 'A st',...). "street": "%s st" % (chr(65 | (i & 14) | ((i & 1) * 32))), } if attrs is not None: user.update(attrs) user['dn'] = "cn=%s,%s" % (user['cn'], self.ou) if opts.skip_attr_regex: match = re.compile(opts.skip_attr_regex).search for k in user.keys(): if match(k): del user[k] self.users.append(user) self.ldb.add(user) return user def setUp(self): super(VLVTests, self).setUp() self.ldb = SamDB(host, credentials=creds, session_info=system_session(lp), lp=lp) self.base_dn = self.ldb.domain_dn() self.ou = "ou=vlv,%s" % self.base_dn if opts.delete_in_setup: try: self.ldb.delete(self.ou, ['tree_delete:1']) except ldb.LdbError, e: print "tried deleting %s, got error %s" % (self.ou, e) self.ldb.add({ "dn": self.ou, "objectclass": "organizationalUnit"}) self.users = [] for i in range(N_ELEMENTS): self.create_user(i, N_ELEMENTS) attrs = self.users[0].keys() self.binary_sorted_keys = ['audio', 'photo', "msTSExpireDate4", 'serialNumber', "displayNamePrintable"] self.numeric_sorted_keys = ['flags', 'accountExpires'] self.timestamp_keys = ['msTSExpireDate4'] self.int64_keys = set(['accountExpires']) self.locale_sorted_keys = [x for x in attrs if x not in (self.binary_sorted_keys + self.numeric_sorted_keys)] # don't try spaces, etc in cn self.delicate_keys = ['cn'] def tearDown(self): super(VLVTests, self).tearDown() if not opts.delete_in_setup: self.ldb.delete(self.ou, ['tree_delete:1']) def get_full_list(self, attr, include_cn=False): """Fetch the whole list sorted on the attribute, using the VLV. This way you get a VLV cookie.""" n_users = len(self.users) sort_control = "server_sort:1:0:%s" % attr half_n = n_users // 2 vlv_search = "vlv:1:%d:%d:%d:0" % (half_n, half_n, half_n + 1) attrs = [attr] if include_cn: attrs.append('cn') res = self.ldb.search(self.ou, scope=ldb.SCOPE_ONELEVEL, attrs=attrs, controls=[sort_control, vlv_search]) if include_cn: full_results = [(x[attr][0], x['cn'][0]) for x in res] else: full_results = [x[attr][0].lower() for x in res] controls = res.controls return full_results, controls, sort_control def get_expected_order(self, attr): """Fetch the whole list sorted on the attribute, using sort only.""" sort_control = "server_sort:1:0:%s" % attr res = self.ldb.search(self.ou, scope=ldb.SCOPE_ONELEVEL, attrs=[attr], controls=[sort_control]) results = [x[attr][0] for x in res] return results def delete_user(self, user): self.ldb.delete(user['dn']) del self.users[self.users.index(user)] def get_gte_tests_and_order(self, attr): expected_order = self.get_expected_order(attr) gte_users = [] if attr in self.delicate_keys: gte_keys = [ '3', 'abc', '¹', 'ŋđ¼³ŧ“«đð', '桑巴', ] elif attr in self.timestamp_keys: gte_keys = [ '18560101010000.0Z', '19140103010000.0Z', '19560101010010.0Z', '19700101000000.0Z', '19991231211234.3Z', '20061111211234.0Z', '20390901041234.0Z', '25560101010000.0Z', ] elif attr not in self.numeric_sorted_keys: gte_keys = [ '3', 'abc', ' ', '!@#!@#!', 'kōkako', '¹', 'ŋđ¼³ŧ“«đð', '\n\t\t', '桑巴', 'zzzz', expected_order[len(expected_order) // 2] + ' tail', ] else: # "numeric" means positive integers # doesn't work with -1, 3.14, ' 3', '9' * 20 gte_keys = ['3', '1' * 10, '1', '9' * 7, '0'] if attr in self.int64_keys: gte_keys += ['3' * 12, '71' * 8] for i, x in enumerate(gte_keys): user = self.create_user(i, N_ELEMENTS, prefix='gte', attrs={attr: x}) gte_users.append(user) gte_order = self.get_expected_order(attr) for user in gte_users: self.delete_user(user) # for sanity's sake expected_order_2 = self.get_expected_order(attr) self.assertEqual(expected_order, expected_order_2) # Map gte tests to indexes in expected order. This will break # if gte_order and expected_order are differently ordered (as # it should). gte_map = {} # index to the first one with each value index_map = {} for i, k in enumerate(expected_order): if k not in index_map: index_map[k] = i keys = [] for k in gte_order: if k in index_map: i = index_map[k] gte_map[k] = i for k in keys: gte_map[k] = i keys = [] else: keys.append(k) for k in keys: gte_map[k] = len(expected_order) if False: print "gte_map:" for k in gte_order: print " %10s => %10s" % (k, gte_map[k]) return gte_order, expected_order, gte_map def assertCorrectResults(self, results, expected_order, offset, before, after): """A helper to calculate offsets correctly and say as much as possible when something goes wrong.""" start = max(offset - before - 1, 0) end = offset + after expected_results = expected_order[start: end] # if it is a tuple with the cn, drop the cn if expected_results and isinstance(expected_results[0], tuple): expected_results = [x[0] for x in expected_results] if expected_results == results: return if expected_order is not None: print "expected order: %s" % expected_order print "offset %d before %d after %d" % (offset, before, after) print "start %d end %d" % (start, end) print "expected: %s" % expected_results print "got : %s" % results self.assertEquals(expected_results, results) def test_server_vlv_with_cookie(self): attrs = [x for x in self.users[0].keys() if x not in ('dn', 'objectclass')] for attr in attrs: expected_order = self.get_expected_order(attr) sort_control = "server_sort:1:0:%s" % attr res = None n = len(self.users) for before in range(0, 9): for after in range(0, 9): for offset in range(max(1, before - 2), min(n - after + 2, n)): if res is None: vlv_search = "vlv:1:%d:%d:%d:0" % (before, after, offset) else: cookie = get_cookie(res.controls, n) vlv_search = ("vlv:1:%d:%d:%d:%s:%s" % (before, after, offset, n, cookie)) res = self.ldb.search(self.ou, scope=ldb.SCOPE_ONELEVEL, attrs=[attr], controls=[sort_control, vlv_search]) results = [x[attr][0] for x in res] self.assertCorrectResults(results, expected_order, offset, before, after) def test_server_vlv_with_cookie_while_adding_and_deleting(self): """What happens if we add or remove items in the middle of the VLV? Nothing. The search and the sort is not repeated, and we only deal with ther objects originally found. """ attrs = ['cn'] + [x for x in self.users[0].keys() if x not in ('dn', 'objectclass')] user_number = 0 iteration = 0 for attr in attrs: full_results, controls, sort_control = \ self.get_full_list(attr, True) original_n = len(self.users) expected_order = full_results random.seed(1) for before in range(0, 3) + [6, 11, 19]: for after in range(0, 3) + [6, 11, 19]: start = max(before - 1, 1) end = max(start + 4, original_n - after + 2) for offset in range(start, end): #if iteration > 2076: # return cookie = get_cookie(controls, original_n) vlv_search = encode_vlv_control(before=before, after=after, offset=offset, n=original_n, cookie=cookie) iteration += 1 res = self.ldb.search(self.ou, scope=ldb.SCOPE_ONELEVEL, attrs=[attr], controls=[sort_control, vlv_search]) controls = res.controls results = [x[attr][0] for x in res] real_offset = max(1, min(offset, len(expected_order))) expected_results = [] skipped = 0 begin_offset = max(real_offset - before - 1, 0) real_before = min(before, real_offset - 1) real_after = min(after, len(expected_order) - real_offset) for x in expected_order[begin_offset:]: if x is not None: expected_results.append(x[0]) if (len(expected_results) == real_before + real_after + 1): break else: skipped += 1 if expected_results != results: print ("attr %s before %d after %d offset %d" % (attr, before, after, offset)) self.assertEquals(expected_results, results) n = len(self.users) if random.random() < 0.1 + (n < 5) * 0.05: if n == 0: i = 0 else: i = random.randrange(n) user = self.create_user(i, n, suffix='-%s' % user_number) user_number += 1 if random.random() < 0.1 + (n > 50) * 0.02 and n: index = random.randrange(n) user = self.users.pop(index) self.ldb.delete(user['dn']) replaced = (user[attr], user['cn']) if replaced in expected_order: i = expected_order.index(replaced) expected_order[i] = None def test_server_vlv_with_cookie_while_changing(self): """What happens if we modify items in the middle of the VLV? The expected behaviour (as found on Windows) is the sort is not repeated, but the changes in attributes are reflected. """ attrs = [x for x in self.users[0].keys() if x not in ('dn', 'objectclass', 'cn')] for attr in attrs: n_users = len(self.users) expected_order = [x.upper() for x in self.get_expected_order(attr)] sort_control = "server_sort:1:0:%s" % attr res = None i = 0 # First we'll fetch the whole list so we know the original # sort order. This is necessary because we don't know how # the server will order equivalent items. We are using the # dn as a key. half_n = n_users // 2 vlv_search = "vlv:1:%d:%d:%d:0" % (half_n, half_n, half_n + 1) res = self.ldb.search(self.ou, scope=ldb.SCOPE_ONELEVEL, attrs=['dn', attr], controls=[sort_control, vlv_search]) results = [x[attr][0].upper() for x in res] #self.assertEquals(expected_order, results) dn_order = [str(x['dn']) for x in res] values = results[:] for before in range(0, 3): for after in range(0, 3): for offset in range(1 + before, n_users - after): cookie = get_cookie(res.controls, len(self.users)) vlv_search = ("vlv:1:%d:%d:%d:%s:%s" % (before, after, offset, len(self.users), cookie)) res = self.ldb.search(self.ou, scope=ldb.SCOPE_ONELEVEL, attrs=['dn', attr], controls=[sort_control, vlv_search]) dn_results = [str(x['dn']) for x in res] dn_expected = dn_order[offset - before - 1: offset + after] self.assertEquals(dn_expected, dn_results) results = [x[attr][0].upper() for x in res] self.assertCorrectResults(results, values, offset, before, after) i += 1 if i % 3 == 2: if (attr in self.locale_sorted_keys or attr in self.binary_sorted_keys): i1 = i % n_users i2 = (i ^ 255) % n_users dn1 = dn_order[i1] dn2 = dn_order[i2] v2 = values[i2] if v2 in self.locale_sorted_keys: v2 += '-%d' % i cn1 = dn1.split(',', 1)[0][3:] cn2 = dn2.split(',', 1)[0][3:] values[i1] = v2 m = ldb.Message() m.dn = ldb.Dn(self.ldb, dn1) m[attr] = ldb.MessageElement(v2, ldb.FLAG_MOD_REPLACE, attr) self.ldb.modify(m) def test_server_vlv_fractions_with_cookie(self): """What happens when the count is set to an arbitrary number? In that case the offset and the count form a fraction, and the VLV should be centred at a point offset/count of the way through. For example, if offset is 3 and count is 6, the VLV should be looking around halfway. The actual algorithm is a bit fiddlier than that, because of the one-basedness of VLV. """ attrs = [x for x in self.users[0].keys() if x not in ('dn', 'objectclass')] n_users = len(self.users) random.seed(4) for attr in attrs: full_results, controls, sort_control = self.get_full_list(attr) self.assertEqual(len(full_results), n_users) for before in range(0, 2): for after in range(0, 2): for denominator in range(1, 20): for offset in range(1, denominator + 3): cookie = get_cookie(controls, len(self.users)) vlv_search = ("vlv:1:%d:%d:%d:%s:%s" % (before, after, offset, denominator, cookie)) try: res = self.ldb.search(self.ou, scope=ldb.SCOPE_ONELEVEL, attrs=[attr], controls=[sort_control, vlv_search]) except ldb.LdbError, e: if offset != 0: raise print ("offset %d denominator %d raised error " "expected error %s\n" "(offset zero is illegal unless " "content count is zero)" % (offset, denominator, e)) continue results = [x[attr][0].lower() for x in res] if denominator == 0: denominator = n_users if offset == 0: offset = denominator elif denominator == 1: # the offset can only be 1, but the 1/1 case # means something special if offset == 1: real_offset = n_users else: real_offset = 1 else: if offset > denominator: offset = denominator real_offset = (1 + int(round((n_users - 1) * (offset - 1) / (denominator - 1.0))) ) self.assertCorrectResults(results, full_results, real_offset, before, after) controls = res.controls if False: for c in list(controls): cstr = str(c) if cstr.startswith('vlv_resp'): bits = cstr.rsplit(':') print ("the answer is %s; we said %d" % (bits[2], real_offset)) break def test_server_vlv_no_cookie(self): attrs = [x for x in self.users[0].keys() if x not in ('dn', 'objectclass')] for attr in attrs: expected_order = self.get_expected_order(attr) sort_control = "server_sort:1:0:%s" % attr for before in range(0, 5): for after in range(0, 7): for offset in range(1 + before, len(self.users) - after): res = self.ldb.search(self.ou, scope=ldb.SCOPE_ONELEVEL, attrs=[attr], controls=[sort_control, "vlv:1:%d:%d:%d:0" % (before, after, offset)]) results = [x[attr][0] for x in res] self.assertCorrectResults(results, expected_order, offset, before, after) def test_server_vlv_gte_with_cookie(self): attrs = [x for x in self.users[0].keys() if x not in ('dn', 'objectclass')] for attr in attrs: gte_order, expected_order, gte_map = \ self.get_gte_tests_and_order(attr) # In case there is some order dependency, disorder tests gte_tests = gte_order[:] random.seed(1) random.shuffle(gte_tests) res = None sort_control = "server_sort:1:0:%s" % attr for before in range(0, 5): for after in range(0, 7): for gte in gte_tests: if res is not None: cookie = get_cookie(res.controls, len(self.users)) else: cookie = None vlv_search = encode_vlv_control(before=before, after=after, gte=gte, cookie=cookie) res = self.ldb.search(self.ou, scope=ldb.SCOPE_ONELEVEL, attrs=[attr], controls=[sort_control, vlv_search]) results = [x[attr][0] for x in res] offset = gte_map.get(gte, len(expected_order)) # here offset is 0-based start = max(offset - before, 0) end = offset + 1 + after expected_results = expected_order[start: end] self.assertEquals(expected_results, results) def test_server_vlv_gte_no_cookie(self): attrs = [x for x in self.users[0].keys() if x not in ('dn', 'objectclass')] iteration = 0 for attr in attrs: gte_order, expected_order, gte_map = \ self.get_gte_tests_and_order(attr) # In case there is some order dependency, disorder tests gte_tests = gte_order[:] random.seed(1) random.shuffle(gte_tests) sort_control = "server_sort:1:0:%s" % attr for before in range(0, 4): for after in range(0, 5): for gte in gte_tests: if attr == 'audio' and 0: import pdb pdb.set_trace() vlv_search = encode_vlv_control(before=before, after=after, gte=gte) res = self.ldb.search(self.ou, scope=ldb.SCOPE_ONELEVEL, attrs=[attr], controls=[sort_control, vlv_search]) results = [x[attr][0] for x in res] # here offset is 0-based offset = gte_map.get(gte, len(expected_order)) start = max(offset - before, 0) end = offset + after + 1 expected_results = expected_order[start: end] iteration += 1 if expected_results != results: middle = expected_order[len(expected_order) // 2] print expected_results, results print middle print expected_order print print ("\nattr %s offset %d before %d " "after %d gte %s" % (attr, offset, before, after, gte)) self.assertEquals(expected_results, results) def test_multiple_searches(self): """The maximum number of concurrent vlv searches per connection is currently set at 3. That means if you open 4 VLV searches the cookie on the first one should fail. """ # Windows has a limit of 10 VLVs where there are low numbers # of objects in each search. attrs = ([x for x in self.users[0].keys() if x not in ('dn', 'objectclass')] * 2)[:12] vlv_cookies = [] for attr in attrs: sort_control = "server_sort:1:0:%s" % attr res = self.ldb.search(self.ou, scope=ldb.SCOPE_ONELEVEL, attrs=[attr], controls=[sort_control, "vlv:1:1:1:1:0"]) cookie = get_cookie(res.controls, len(self.users)) vlv_cookies.append(cookie) time.sleep(0.2) # now this one should fail self.assertRaises(ldb.LdbError, self.ldb.search, self.ou, scope=ldb.SCOPE_ONELEVEL, attrs=[attr], controls=[sort_control, "vlv:1:1:1:1:0:%s" % vlv_cookies[0]]) # and this one should succeed res = self.ldb.search(self.ou, scope=ldb.SCOPE_ONELEVEL, attrs=[attr], controls=[sort_control, "vlv:1:1:1:1:0:%s" % vlv_cookies[-1]]) # this one should fail because it is a new connection and # doesn't share cookies new_ldb = SamDB(host, credentials=creds, session_info=system_session(lp), lp=lp) self.assertRaises(ldb.LdbError, new_ldb.search, self.ou, scope=ldb.SCOPE_ONELEVEL, attrs=[attr], controls=[sort_control, "vlv:1:1:1:1:0:%s" % vlv_cookies[-1]]) # but now without the critical flag it just does no VLV. new_ldb.search(self.ou, scope=ldb.SCOPE_ONELEVEL, attrs=[attr], controls=[sort_control, "vlv:0:1:1:1:0:%s" % vlv_cookies[-1]]) if "://" not in host: if os.path.isfile(host): host = "tdb://%s" % host else: host = "ldap://%s" % host TestProgram(module=__name__, opts=subunitopts)