samba-tool visualize ntdsconn: add --importldif option

This visualizes the NTDSConnections in an LDIF file exported via
`samba_kcc --exportldif`. This functionality is already available in a
roundabout way -- you can use `samba_kcc --import_ldif`, and use the
DB that generates. This just shortens the process.

The ldif import/export feature is useful for analysing AD networks
offsite without exposing too much sensitive data.

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

2 files changed, 41 insertions, 3 deletions

diff --git a/python/samba/netcmd/visualize.py b/python/samba/netcmd/visualize.py
index 311476a60ef..e66030d3107 100644
--- a/python/samba/netcmd/visualize.py
+++ b/python/samba/netcmd/visualize.py
@@ -33,7 +33,7 @@ from samba.graph import distance_matrix, COLOUR_SETS
 from ldb import SCOPE_BASE, SCOPE_SUBTREE, LdbError
 import time
 import re
-from samba.kcc import KCC
+from samba.kcc import KCC, ldif_import_export
 from samba.kcc.kcc_utils import KCCError
 from samba.compat import text_type
 
@@ -400,13 +400,31 @@ class NTDSConn(object):
 
 class cmd_ntdsconn(GraphCommand):
     "Draw the NTDSConnection graph"
+    takes_options = COMMON_OPTIONS + [
+        Option("--importldif", help="graph from samba_kcc generated ldif",
+               default=None),
+    ]
+
+    def import_ldif_db(self, ldif, lp):
+        d = tempfile.mkdtemp(prefix='samba-tool-visualise')
+        fn = os.path.join(d, 'imported.ldb')
+        self._tmp_fn_to_delete = fn
+        samdb = ldif_import_export.ldif_to_samdb(fn, lp, ldif)
+        return fn
+
     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=None, format=None):
+            utf8=None, format=None, importldif=None):
+
         lp = sambaopts.get_loadparm()
-        creds = credopts.get_credentials(lp, fallback_machine=True)
+        if importldif is None:
+            creds = credopts.get_credentials(lp, fallback_machine=True)
+        else:
+            creds = None
+            H = self.import_ldif_db(importldif, lp)
+
         local_kcc, dsas = self.get_kcc_and_dsas(H, lp, creds)
         local_dsa_dn = local_kcc.my_dsa_dnstr.split(',', 1)[1]
         vertices = set()
@@ -449,6 +467,10 @@ class cmd_ntdsconn(GraphCommand):
                 attested_edges.append((msg['fromServer'][0],
                                        dest_dn, ntds_dn))
 
+        if importldif and H == self._tmp_fn_to_delete:
+            os.remove(H)
+            os.rmdir(os.path.dirname(H))
+
         # now we overlay all the graphs and generate styles accordingly
         edges = {}
         for src, dest, attester in attested_edges:
diff --git a/python/samba/tests/samba_tool/visualize.py b/python/samba/tests/samba_tool/visualize.py
index c015abf3b2d..1afb210b7bd 100644
--- a/python/samba/tests/samba_tool/visualize.py
+++ b/python/samba/tests/samba_tool/visualize.py
@@ -122,6 +122,22 @@ class SambaToolVisualizeLdif(SambaToolCmdTest):
 
             self.assertStringsEqual(monochrome, uncoloured, strip=True)
 
+    def test_import_ldif(self):
+        """Make sure the samba-tool visualize --importldif option gives the
+        same output as using the externally generated db from the same
+        LDIF."""
+        result, s1, err = self.runsubcmd("visualize", "ntdsconn",
+                                         '-H', self.dburl,
+                                         '--color=no', '-S')
+        self.assertCmdSuccess(result, s1, err)
+
+        result, s2, err = self.runsubcmd("visualize", "ntdsconn",
+                                         '--importldif', MULTISITE_LDIF,
+                                         '--color=no', '-S')
+        self.assertCmdSuccess(result, s2, err)
+
+        self.assertStringsEqual(s1, s2)
+
     def test_output_file(self):
         """Check that writing to a file works, with and without
         --color=auto."""