sambatool visualize: add up-to-dateness visualization

Or more accurately, out-of-dateness visualization, which shows how far
each DCs is from every other using the difference in the up-to-dateness
vectors.

An example usage is

samba-tool visualize uptodateness -r -S -H ldap://somewhere \
      -UAdministrator --color=auto --partition=DOMAIN

Signed-off-by: Douglas Bagnall <douglas.bagnall@catalyst.net.nz>
Reviewed-by: Andrew Bartlett <abartlet@samba.org>

2 files changed, 532 insertions, 4 deletions

diff --git a/python/samba/netcmd/visualize.py b/python/samba/netcmd/visualize.py
index bfd7d3bf341..a24962ea58a 100644
--- a/python/samba/netcmd/visualize.py
+++ b/python/samba/netcmd/visualize.py
@@ -25,12 +25,14 @@ from collections import defaultdict
 import subprocess
 
 import tempfile
-import samba
 import samba.getopt as options
+from samba import dsdb
+from samba import nttime2unix
 from samba.netcmd import Command, SuperCommand, CommandError, Option
 from samba.samdb import SamDB
 from samba.graph import dot_graph
 from samba.graph import distance_matrix, COLOUR_SETS
+from samba.graph import full_matrix
 from ldb import SCOPE_BASE, SCOPE_SUBTREE, LdbError
 import time
 import re
@@ -672,6 +674,140 @@ class cmd_ntdsconn(GraphCommand):
             self.write(s, output)
 
 
+class cmd_uptodateness(GraphCommand):
+    """visualize uptodateness vectors"""
+
+    takes_options = COMMON_OPTIONS + [
+        Option("-p", "--partition", help="restrict to this partition",
+               default=None),
+        Option("--max-digits", default=3, type=int,
+               help="display this many digits of out-of-date-ness"),
+    ]
+
+    def get_utdv(self, samdb, dn):
+        """This finds the uptodateness vector in the database."""
+        cursors = []
+        config_dn = samdb.get_config_basedn()
+        for c in dsdb._dsdb_load_udv_v2(samdb, dn):
+            inv_id = str(c.source_dsa_invocation_id)
+            res = samdb.search(base=config_dn,
+                               expression=("(&(invocationId=%s)"
+                                           "(objectClass=nTDSDSA))" % inv_id),
+                               attrs=["distinguishedName", "invocationId"])
+            settings_dn = res[0]["distinguishedName"][0]
+            prefix, dsa_dn = settings_dn.split(',', 1)
+            if prefix != 'CN=NTDS Settings':
+                raise CommandError("Expected NTDS Settings DN, got %s" %
+                                   settings_dn)
+
+            cursors.append((dsa_dn,
+                            inv_id,
+                            int(c.highest_usn),
+                            nttime2unix(c.last_sync_success)))
+        return cursors
+
+    def get_own_cursor(self, samdb):
+            res = samdb.search(base="",
+                               scope=SCOPE_BASE,
+                               attrs=["highestCommittedUSN"])
+            usn = int(res[0]["highestCommittedUSN"][0])
+            now = int(time.time())
+            return (usn, now)
+
+    def run(self, H=None, output=None, shorten_names=False,
+            key=True, talk_to_remote=False,
+            sambaopts=None, credopts=None, versionopts=None,
+            color=None, color_scheme=None,
+            utf8=False, format=None, importldif=None,
+            xdot=False, partition=None, max_digits=3):
+        if not talk_to_remote:
+            print("this won't work without talking to the remote servers "
+                  "(use -r)", file=self.outf)
+            return
+
+        # We use the KCC libraries in readonly mode to get the
+        # replication graph.
+        lp = sambaopts.get_loadparm()
+        creds = credopts.get_credentials(lp, fallback_machine=True)
+        local_kcc, dsas = self.get_kcc_and_dsas(H, lp, creds)
+        self.samdb = local_kcc.samdb
+        partition = get_partition(self.samdb, partition)
+
+        short_partitions, long_partitions = get_partition_maps(self.samdb)
+        color_scheme = self.calc_distance_color_scheme(color,
+                                                       color_scheme,
+                                                       output)
+
+        for part_name, part_dn in short_partitions.items():
+            if partition not in (part_dn, None):
+                continue  # we aren't doing this partition
+
+            cursors = self.get_utdv(self.samdb, part_dn)
+
+            # we talk to each remote and make a matrix of the vectors
+            # -- for each partition
+            # normalise by oldest
+            utdv_edges = {}
+            for dsa_dn in dsas:
+                res = local_kcc.samdb.search(dsa_dn,
+                                             scope=SCOPE_BASE,
+                                             attrs=["dNSHostName"])
+                ldap_url = "ldap://%s" % res[0]["dNSHostName"][0]
+                try:
+                    samdb = self.get_db(ldap_url, sambaopts, credopts)
+                    cursors = self.get_utdv(samdb, part_dn)
+                    own_usn, own_time = self.get_own_cursor(samdb)
+                    remotes = {dsa_dn: own_usn}
+                    for dn, guid, usn, t in cursors:
+                        remotes[dn] = usn
+                except LdbError as e:
+                    print("Could not contact %s (%s)" % (ldap_url, e),
+                          file=sys.stderr)
+                    continue
+                utdv_edges[dsa_dn] = remotes
+
+            distances = {}
+            max_distance = 0
+            for dn1 in dsas:
+                try:
+                    peak = utdv_edges[dn1][dn1]
+                except KeyError as e:
+                    peak = 0
+                d = {}
+                distances[dn1] = d
+                for dn2 in dsas:
+                    if dn2 in utdv_edges:
+                        if dn1 in utdv_edges[dn2]:
+                            dist = peak - utdv_edges[dn2][dn1]
+                            d[dn2] = dist
+                            if dist > max_distance:
+                                max_distance = dist
+                        else:
+                            print("Missing dn %s from UTD vector" % dn1,
+                                  file=sys.stderr)
+                    else:
+                        print("missing dn %s from UTD vector list" % dn2,
+                              file=sys.stderr)
+
+            digits = min(max_digits, len(str(max_distance)))
+            if digits < 1:
+                digits = 1
+            c_scale = 10 ** digits
+
+            s = full_matrix(distances,
+                            utf8=utf8,
+                            colour=color_scheme,
+                            shorten_names=shorten_names,
+                            generate_key=key,
+                            grouping_function=get_dnstr_site,
+                            colour_scale=c_scale,
+                            digits=digits,
+                            ylabel='DC',
+                            xlabel='out-of-date-ness')
+
+            self.write('\n%s\n\n%s' % (part_name, s), output)
+
+
 class cmd_visualize(SuperCommand):
     """Produces graphical representations of Samba network state"""
     subcommands = {}
diff --git a/python/samba/tests/samba_tool/visualize_drs.py b/python/samba/tests/samba_tool/visualize_drs.py
index 7da0a4b1083..42facacd977 100644
--- a/python/samba/tests/samba_tool/visualize_drs.py
+++ b/python/samba/tests/samba_tool/visualize_drs.py
@@ -1,3 +1,4 @@
+# -*- coding: utf-8 -*-
 # Originally based on tests for samba.kcc.ldif_import_export.
 # Copyright (C) Andrew Bartlett 2015, 2018
 #
@@ -16,15 +17,22 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
-
 """Tests for samba-tool visualize using the vampire DC and promoted DC
-environments. We can't assert much about what state they are in, so we
-mainly check for cmmand failure.
+environments. For most tests we assume we can't assert much about what
+state they are in, so we mainly check for command failure, but for
+others we try to grasp control of replication and make more specific
+assertions.
 """
 
+from __future__ import print_function
 import os
+import re
+import random
+import subprocess
 from samba.tests.samba_tool.base import SambaToolCmdTest
 
+VERBOSE = False
+
 ENV_DSAS = {
     'promoted_dc': ['CN=PROMOTEDVDC,CN=Servers,CN=Default-First-Site-Name,CN=Sites,CN=Configuration,DC=samba,DC=example,DC=com',
                     'CN=LOCALDC,CN=Servers,CN=Default-First-Site-Name,CN=Sites,CN=Configuration,DC=samba,DC=example,DC=com'],
@@ -33,6 +41,59 @@ ENV_DSAS = {
 }
 
 
+def set_auto_replication(dc, allow):
+    credstring = '-U%s%%%s' % (os.environ["USERNAME"], os.environ["PASSWORD"])
+    on_or_off = '-' if allow else '+'
+
+    for opt in ['DISABLE_INBOUND_REPL',
+                'DISABLE_OUTBOUND_REPL']:
+        cmd = ['bin/samba-tool',
+               'drs', 'options',
+               credstring, dc,
+               "--dsa-option=%s%s" % (on_or_off, opt)]
+
+        subprocess.check_call(cmd)
+
+
+def force_replication(src, dest, base):
+    credstring = '-U%s%%%s' % (os.environ["USERNAME"], os.environ["PASSWORD"])
+    cmd = ['bin/samba-tool',
+           'drs', 'replicate',
+           dest, src, base,
+           credstring,
+           '--sync-forced']
+
+    subprocess.check_call(cmd)
+
+
+def get_utf8_matrix(s):
+    # parse the graphical table *just* well enough for our tests
+    # decolourise first
+    s = re.sub("\033" r"\[[^m]+m", '', s)
+    lines = s.split('\n')
+    # matrix rows have '·' on the diagonal
+    rows = [x.strip().replace('·', '0') for x in lines if '·' in x]
+    names = []
+    values = []
+    for r in rows:
+        parts = r.rsplit(None, len(rows))
+        k, v = parts[0], parts[1:]
+        # we want the FOO in 'CN=FOO+' or 'CN=FOO,CN=x,DC=...'
+        k = re.match(r'cn=([^+,]+)', k.lower()).group(1)
+        names.append(k)
+        if len(v) == 1:  # this is a single-digit matrix, no spaces
+            v = list(v[0])
+        values.append([int(x) if x.isdigit() else 1e999 for x in v])
+
+    d = {}
+    for n1, row in zip(names, values):
+        d[n1] = {}
+        for n2, v in zip(names, row):
+            d[n1][n2] = v
+
+    return d
+
+
 class SambaToolVisualizeDrsTest(SambaToolCmdTest):
     def setUp(self):
         super(SambaToolVisualizeDrsTest, self).setUp()
@@ -64,6 +125,337 @@ class SambaToolVisualizeDrsTest(SambaToolCmdTest):
                                             '--color=no', '-S')
         self.assertCmdSuccess(result, out, err)
 
+    def test_uptodateness_all_partitions(self):
+        creds = "%s%%%s" % (os.environ["USERNAME"], os.environ["PASSWORD"])
+        dc1 = os.environ["SERVER"]
+        dc2 = os.environ["DC_SERVER"]
+        # We will check that the visualisation works for the two
+        # stopped DCs, but we can't make assertions that the output
+        # will be the same because there may be replication between
+        # the two calls. Stopping the replication on these ones is not
+        # enough because there are other DCs about.
+        (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                            "-r",
+                                            '-H', "ldap://%s" % dc1,
+                                            '-U', creds,
+                                            '--color=no', '-S')
+        self.assertCmdSuccess(result, out, err)
+
+        (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                            "-r",
+                                            '-H', "ldap://%s" % dc2,
+                                            '-U', creds,
+                                            '--color=no', '-S')
+        self.assertCmdSuccess(result, out, err)
+
+    def test_uptodateness_partitions(self):
+        creds = "%s%%%s" % (os.environ["USERNAME"], os.environ["PASSWORD"])
+        dc1 = os.environ["SERVER"]
+        for part in ["CONFIGURATION",
+                     "SCHEMA",
+                     "DNSDOMAIN",
+                     "DNSFOREST"]:
+
+            (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                                "-r",
+                                                '-H', "ldap://%s" % dc1,
+                                                '-U', creds,
+                                                '--color=no', '-S',
+                                                '--partition', part)
+            self.assertCmdSuccess(result, out, err)
+
+    def assert_matrix_validity(self, matrix, dcs=()):
+        for dc in dcs:
+            self.assertIn(dc, matrix)
+        for k, row in matrix.items():
+            self.assertEqual(row[k], 0)
+
+    def test_uptodateness_stop_replication_domain(self):
+        creds = "%s%%%s" % (os.environ["USERNAME"], os.environ["PASSWORD"])
+        dc1 = os.environ["SERVER"]
+        dc2 = os.environ["DC_SERVER"]
+        self.addCleanup(set_auto_replication, dc1, True)
+        self.addCleanup(set_auto_replication, dc2, True)
+
+        def display(heading, out):
+            if VERBOSE:
+                print("========", heading, "=========")
+                print(out)
+
+        samdb1 = self.getSamDB("-H", "ldap://%s" % dc1, "-U", creds)
+        samdb2 = self.getSamDB("-H", "ldap://%s" % dc2, "-U", creds)
+
+        domain_dn = samdb1.domain_dn()
+        self.assertTrue(domain_dn == samdb2.domain_dn(),
+                        "We expected the same domain_dn across DCs")
+
+        ou1 = "OU=dc1.%x,%s" % (random.randrange(1 << 64), domain_dn)
+        ou2 = "OU=dc2.%x,%s" % (random.randrange(1 << 64), domain_dn)
+        samdb1.add({
+            "dn": ou1,
+            "objectclass": "organizationalUnit"
+        })
+        samdb2.add({
+            "dn": ou2,
+            "objectclass": "organizationalUnit"
+        })
+
+        set_auto_replication(dc1, False)
+        (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                            "-r",
+                                            '-H', "ldap://%s" % dc1,
+                                            '-U', creds,
+                                            '--color=yes',
+                                            '--utf8', '-S',
+                                            '--partition', 'DOMAIN')
+        display("dc1 replication is now off", out)
+        self.assertCmdSuccess(result, out, err)
+        matrix = get_utf8_matrix(out)
+        self.assert_matrix_validity(matrix, [dc1, dc2])
+
+        force_replication(dc2, dc1, domain_dn)
+        (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                            "-r",
+                                            '-H', "ldap://%s" % dc1,
+                                            '-U', creds,
+                                            '--color=yes',
+                                            '--utf8', '-S',
+                                            '--partition', 'DOMAIN')
+        display("forced replication %s -> %s" % (dc2, dc1), out)
+        self.assertCmdSuccess(result, out, err)
+        matrix = get_utf8_matrix(out)
+        self.assert_matrix_validity(matrix, [dc1, dc2])
+        self.assertEqual(matrix[dc1][dc2], 0)
+
+        force_replication(dc1, dc2, domain_dn)
+        (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                            "-r",
+                                            '-H', "ldap://%s" % dc1,
+                                            '-U', creds,
+                                            '--color=yes',
+                                            '--utf8', '-S',
+                                            '--partition', 'DOMAIN')
+        display("forced replication %s -> %s" % (dc2, dc1), out)
+        self.assertCmdSuccess(result, out, err)
+        matrix = get_utf8_matrix(out)
+        self.assert_matrix_validity(matrix, [dc1, dc2])
+        self.assertEqual(matrix[dc2][dc1], 0)
+
+        dn1 = 'cn=u1.%%d,%s' % (ou1)
+        dn2 = 'cn=u2.%%d,%s' % (ou2)
+
+        for i in range(10):
+            samdb1.add({
+                "dn": dn1 % i,
+                "objectclass": "user"
+            })
+
+        (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                            "-r",
+                                            '-H', "ldap://%s" % dc1,
+                                            '-U', creds,
+                                            '--color=yes',
+                                            '--utf8', '-S',
+                                            '--partition', 'DOMAIN')
+        display("added 10 users on %s" % dc1, out)
+        self.assertCmdSuccess(result, out, err)
+        matrix = get_utf8_matrix(out)
+        self.assert_matrix_validity(matrix, [dc1, dc2])
+        # dc2's view of dc1 should now be 10 changes out of date
+        self.assertEqual(matrix[dc2][dc1], 10)
+
+        for i in range(10):
+            samdb2.add({
+                "dn": dn2 % i,
+                "objectclass": "user"
+            })
+
+        (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                            "-r",
+                                            '-H', "ldap://%s" % dc1,
+                                            '-U', creds,
+                                            '--color=yes',
+                                            '--utf8', '-S',
+                                            '--partition', 'DOMAIN')
+        display("added 10 users on %s" % dc2, out)
+        self.assertCmdSuccess(result, out, err)
+        matrix = get_utf8_matrix(out)
+        self.assert_matrix_validity(matrix, [dc1, dc2])
+        # dc1's view of dc2 is probably 11 changes out of date
+        self.assertGreaterEqual(matrix[dc1][dc2], 10)
+
+        for i in range(10, 101):
+            samdb1.add({
+                "dn": dn1 % i,
+                "objectclass": "user"
+            })
+            samdb2.add({
+                "dn": dn2 % i,
+                "objectclass": "user"
+            })
+
+        (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                            "-r",
+                                            '-H', "ldap://%s" % dc1,
+                                            '-U', creds,
+                                            '--color=yes',
+                                            '--utf8', '-S',
+                                            '--partition', 'DOMAIN')
+        display("added 91 users on both", out)
+        self.assertCmdSuccess(result, out, err)
+        matrix = get_utf8_matrix(out)
+        self.assert_matrix_validity(matrix, [dc1, dc2])
+        # the difference here should be ~101.
+        self.assertGreaterEqual(matrix[dc1][dc2], 100)
+        self.assertGreaterEqual(matrix[dc2][dc1], 100)
+
+        (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                            "-r",
+                                            '-H', "ldap://%s" % dc1,
+                                            '-U', creds,
+                                            '--color=yes',
+                                            '--utf8', '-S',
+                                            '--partition', 'DOMAIN',
+                                            '--max-digits', '2')
+        display("with --max-digits 2", out)
+        self.assertCmdSuccess(result, out, err)
+        matrix = get_utf8_matrix(out)
+        self.assert_matrix_validity(matrix, [dc1, dc2])
+        # visualising with 2 digits mean these overflow into infinity
+        self.assertGreaterEqual(matrix[dc1][dc2], 1e99)
+        self.assertGreaterEqual(matrix[dc2][dc1], 1e99)
+
+        (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                            "-r",
+                                            '-H', "ldap://%s" % dc1,
+                                            '-U', creds,
+                                            '--color=yes',
+                                            '--utf8', '-S',
+                                            '--partition', 'DOMAIN',
+                                            '--max-digits', '1')
+        display("with --max-digits 1", out)
+        self.assertCmdSuccess(result, out, err)
+        matrix = get_utf8_matrix(out)
+        self.assert_matrix_validity(matrix, [dc1, dc2])
+        # visualising with 1 digit means these overflow into infinity
+        self.assertGreaterEqual(matrix[dc1][dc2], 1e99)
+        self.assertGreaterEqual(matrix[dc2][dc1], 1e99)
+
+        force_replication(dc2, dc1, samdb1.domain_dn())
+        (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                            "-r",
+                                            '-H', "ldap://%s" % dc1,
+                                            '-U', creds,
+                                            '--color=yes',
+                                            '--utf8', '-S',
+                                            '--partition', 'DOMAIN')
+
+        display("forced replication %s -> %s" % (dc2, dc1), out)
+        self.assertCmdSuccess(result, out, err)
+        matrix = get_utf8_matrix(out)
+        self.assert_matrix_validity(matrix, [dc1, dc2])
+        self.assertEqual(matrix[dc1][dc2], 0)
+
+        force_replication(dc1, dc2, samdb2.domain_dn())
+        (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                            "-r",
+                                            '-H', "ldap://%s" % dc1,
+                                            '-U', creds,
+                                            '--color=yes',
+                                            '--utf8', '-S',
+                                            '--partition', 'DOMAIN')
+
+        display("forced replication %s -> %s" % (dc1, dc2), out)
+        self.assertCmdSuccess(result, out, err)
+        matrix = get_utf8_matrix(out)
+        self.assert_matrix_validity(matrix, [dc1, dc2])
+        self.assertEqual(matrix[dc2][dc1], 0)
+
+        samdb1.delete(ou1, ['tree_delete:1'])
+        samdb2.delete(ou2, ['tree_delete:1'])
+
+        (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                            "-r",
+                                            '-H', "ldap://%s" % dc1,
+                                            '-U', creds,
+                                            '--color=yes',
+                                            '--utf8', '-S',
+                                            '--partition', 'DOMAIN')
+        display("tree delete both ous on %s" % (dc1,), out)
+        self.assertCmdSuccess(result, out, err)
+        matrix = get_utf8_matrix(out)
+        self.assert_matrix_validity(matrix, [dc1, dc2])
+        self.assertGreaterEqual(matrix[dc1][dc2], 100)
+        self.assertGreaterEqual(matrix[dc2][dc1], 100)
+
+        set_auto_replication(dc1, True)
+        (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                            "-r",
+                                            '-H', "ldap://%s" % dc1,
+                                            '-U', creds,
+                                            '--color=yes',
+                                            '--utf8', '-S',
+                                            '--partition', 'DOMAIN')
+        display("replication is now on", out)
+        self.assertCmdSuccess(result, out, err)
+        matrix = get_utf8_matrix(out)
+        self.assert_matrix_validity(matrix, [dc1, dc2])
+        # We can't assert actual values after this because
+        # auto-replication is on and things will change underneath us.
+
+        (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                            "-r",
+                                            '-H', "ldap://%s" % dc2,
+                                            '-U', creds,
+                                            '--color=yes',
+                                            '--utf8', '-S',
+                                            '--partition', 'DOMAIN')
+
+        display("%s's view" % dc2, out)
+        self.assertCmdSuccess(result, out, err)
+        matrix = get_utf8_matrix(out)
+        self.assert_matrix_validity(matrix, [dc1, dc2])
+
+        force_replication(dc1, dc2, samdb2.domain_dn())
+        (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                            "-r",
+                                            '-H', "ldap://%s" % dc1,
+                                            '-U', creds,
+                                            '--color=yes',
+                                            '--utf8', '-S',
+                                            '--partition', 'DOMAIN')
+
+        display("forced replication %s -> %s" % (dc1, dc2), out)
+        self.assertCmdSuccess(result, out, err)
+        matrix = get_utf8_matrix(out)
+        self.assert_matrix_validity(matrix, [dc1, dc2])
+
+        force_replication(dc2, dc1, samdb2.domain_dn())
+        (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                            "-r",
+                                            '-H', "ldap://%s" % dc1,
+                                            '-U', creds,
+                                            '--color=yes',
+                                            '--utf8', '-S',
+                                            '--partition', 'DOMAIN')
+        display("forced replication %s -> %s" % (dc2, dc1), out)
+        self.assertCmdSuccess(result, out, err)
+        matrix = get_utf8_matrix(out)
+        self.assert_matrix_validity(matrix, [dc1, dc2])
+
+        (result, out, err) = self.runsubcmd("visualize", "uptodateness",
+                                            "-r",
+                                            '-H', "ldap://%s" % dc2,
+                                            '-U', creds,
+                                            '--color=yes',
+                                            '--utf8', '-S',
+                                            '--partition', 'DOMAIN')
+        display("%s's view" % dc2, out)
+
+        self.assertCmdSuccess(result, out, err)
+        matrix = get_utf8_matrix(out)
+        self.assert_matrix_validity(matrix, [dc1, dc2])
+
     def test_reps_remote(self):
         server = "ldap://%s" % os.environ["SERVER"]
         creds = "%s%%%s" % (os.environ["USERNAME"], os.environ["PASSWORD"])