KCC: improve samba.kcc.graph.get_spanning_tree_edges() docstring
[nivanova/samba-autobuild/.git] / python / samba / kcc / graph.py
1 # Graph functions used by KCC intersite
2 #
3 # Copyright (C) Dave Craft 2011
4 # Copyright (C) Andrew Bartlett 2015
5 #
6 # Andrew Bartlett's alleged work performed by his underlings Douglas
7 # Bagnall and Garming Sam.
8 #
9 # This program is free software; you can redistribute it and/or modify
10 # it under the terms of the GNU General Public License as published by
11 # the Free Software Foundation; either version 3 of the License, or
12 # (at your option) any later version.
13 #
14 # This program is distributed in the hope that it will be useful,
15 # but WITHOUT ANY WARRANTY; without even the implied warranty of
16 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
17 # GNU General Public License for more details.
18 #
19 # You should have received a copy of the GNU General Public License
20 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
21
22 import itertools
23 import heapq
24
25 from samba.kcc.graph_utils import write_dot_file, verify_and_dot, verify_graph
26 from samba.ndr import ndr_pack
27 from samba.dcerpc import misc
28
29 from samba.kcc.debug import DEBUG, DEBUG_FN
30
31 MAX_DWORD = 2 ** 32 - 1
32
33
34 class ReplInfo(object):
35     """Represents information about replication
36
37     NTDSConnections use one representation a replication schedule, and
38     graph vertices use another. This is the Vertex one.
39
40     """
41     def __init__(self):
42         self.cost = 0
43         self.interval = 0
44         self.options = 0
45         self.schedule = None
46
47
48 def total_schedule(schedule):
49     """Return the total number of 15 minute windows in which the schedule
50     is set to replicate in a week. If the schedule is None it is
51     assumed that the replication will happen in every 15 minute
52     window.
53
54     This is essentially a bit population count.
55     """
56
57     if schedule is None:
58         return 84 * 8  # 84 bytes = 84 * 8 bits
59
60     total = 0
61     for byte in schedule:
62         while byte != 0:
63             total += byte & 1
64             byte >>= 1
65     return total
66
67
68 def convert_schedule_to_repltimes(schedule):
69     """Convert NTDS Connection schedule to replTime schedule.
70
71     Schedule defined in  MS-ADTS 6.1.4.5.2
72     ReplTimes defined in MS-DRSR 5.164.
73
74     "Schedule" has 168 bytes but only the lower nibble of each is
75     significant. There is one byte per hour. Bit 3 (0x08) represents
76     the first 15 minutes of the hour and bit 0 (0x01) represents the
77     last 15 minutes. The first byte presumably covers 12am - 1am
78     Sunday, though the spec doesn't define the start of a week.
79
80     "ReplTimes" has 84 bytes which are the 168 lower nibbles of
81     "Schedule" packed together. Thus each byte covers 2 hours. Bits 7
82     (i.e. 0x80) is the first 15 minutes and bit 0 is the last. The
83     first byte covers Sunday 12am - 2am (per spec).
84
85     Here we pack two elements of the NTDS Connection schedule slots
86     into one element of the replTimes list.
87
88     If no schedule appears in NTDS Connection then a default of 0x11
89     is set in each replTimes slot as per behaviour noted in a Windows
90     DC. That default would cause replication within the last 15
91     minutes of each hour.
92     """
93     if schedule is None or schedule.dataArray[0] is None:
94         return [0x11] * 84
95
96     times = []
97     data = schedule.dataArray[0].slots
98
99     for i in range(84):
100         times.append((data[i * 2] & 0xF) << 4 | (data[i * 2 + 1] & 0xF))
101
102     return times
103
104
105 # Returns true if schedule intersect
106 def combine_repl_info(info_a, info_b, info_c):
107     """Set a replInfo to be the intersection of two others
108
109     If there is any overlap in the replication times specified by the
110     first two parameters, the third replInfo parameter is set up with
111     that overlap, and True is returned. If there is no overlap, the
112     third parameter is unchanged and False is returned.
113
114     :param info_a: An input replInfo object
115     :param info_b: An input replInfo object
116     :param info_c: The result replInfo, set to the intersection of A and B
117                    if the intersection is non-empty.
118     :return: True if info_c allows any replication at all, otherwise False
119     """
120     info_c.interval = max(info_a.interval, info_b.interval)
121     info_c.options = info_a.options & info_b.options
122
123     if info_a.schedule is None:
124         info_a.schedule = [0xFF] * 84
125     if info_b.schedule is None:
126         info_b.schedule = [0xFF] * 84
127
128     new_info = [a & b for a, b in zip(info_a.schedule, info_b.schedule)]
129
130     if not any(new_info):
131         return False
132
133     info_c.schedule = new_info
134
135     # Truncate to MAX_DWORD
136     info_c.cost = info_a.cost + info_b.cost
137     if info_c.cost > MAX_DWORD:
138         info_c.cost = MAX_DWORD
139
140     return True
141
142
143 def get_spanning_tree_edges(graph, my_site, label=None, verify=False,
144                             dot_file_dir=None):
145     """Find edges for the intersite graph
146
147     From MS-ADTS 6.2.2.3.4.4
148
149     :param graph: a kcc.kcc_utils.Graph object
150     :param my_site: the topology generator's site
151     :param label: a label for use in dot files and verification
152     :param verify: if True, try to verify that graph properties are correct
153     :param dot_file_dir: if not None, write Graphviz dot files here
154     """
155     # Phase 1: Run Dijkstra's to get a list of internal edges, which are
156     # just the shortest-paths connecting colored vertices
157
158     internal_edges = set()
159
160     for e_set in graph.edge_set:
161         edgeType = None
162         for v in graph.vertices:
163             v.edges = []
164
165         # All con_type in an edge set is the same
166         for e in e_set.edges:
167             edgeType = e.con_type
168             for v in e.vertices:
169                 v.edges.append(e)
170
171         if verify or dot_file_dir is not None:
172             graph_edges = [(a.site.site_dnstr, b.site.site_dnstr)
173                            for a, b in
174                            itertools.chain(
175                                *(itertools.combinations(edge.vertices, 2)
176                                  for edge in e_set.edges))]
177             graph_nodes = [v.site.site_dnstr for v in graph.vertices]
178
179             if dot_file_dir is not None:
180                 write_dot_file('edgeset_%s' % (edgeType,), graph_edges,
181                                vertices=graph_nodes, label=label)
182
183             if verify:
184                 verify_graph('spanning tree edge set %s' % edgeType,
185                              graph_edges, vertices=graph_nodes,
186                              properties=('complete', 'connected'),
187                              debug=DEBUG)
188
189         # Run dijkstra's algorithm with just the red vertices as seeds
190         # Seed from the full replicas
191         dijkstra(graph, edgeType, False)
192
193         # Process edge set
194         process_edge_set(graph, e_set, internal_edges)
195
196         # Run dijkstra's algorithm with red and black vertices as the seeds
197         # Seed from both full and partial replicas
198         dijkstra(graph, edgeType, True)
199
200         # Process edge set
201         process_edge_set(graph, e_set, internal_edges)
202
203     # All vertices have root/component as itself
204     setup_vertices(graph)
205     process_edge_set(graph, None, internal_edges)
206
207     if verify or dot_file_dir is not None:
208         graph_edges = [(e.v1.site.site_dnstr, e.v2.site.site_dnstr)
209                        for e in internal_edges]
210         graph_nodes = [v.site.site_dnstr for v in graph.vertices]
211         verify_properties = ('multi_edge_forest',)
212         verify_and_dot('prekruskal', graph_edges, graph_nodes, label=label,
213                        properties=verify_properties, debug=DEBUG,
214                        verify=verify, dot_file_dir=dot_file_dir)
215
216     # Phase 2: Run Kruskal's on the internal edges
217     output_edges, components = kruskal(graph, internal_edges)
218
219     # This recalculates the cost for the path connecting the
220     # closest red vertex. Ignoring types is fine because NO
221     # suboptimal edge should exist in the graph
222     dijkstra(graph, "EDGE_TYPE_ALL", False)  # TODO rename
223     # Phase 3: Process the output
224     for v in graph.vertices:
225         if v.is_red():
226             v.dist_to_red = 0
227         else:
228             v.dist_to_red = v.repl_info.cost
229
230     if verify or dot_file_dir is not None:
231         graph_edges = [(e.v1.site.site_dnstr, e.v2.site.site_dnstr)
232                        for e in internal_edges]
233         graph_nodes = [v.site.site_dnstr for v in graph.vertices]
234         verify_properties = ('multi_edge_forest',)
235         verify_and_dot('postkruskal', graph_edges, graph_nodes,
236                        label=label, properties=verify_properties,
237                        debug=DEBUG, verify=verify,
238                        dot_file_dir=dot_file_dir)
239
240     # Ensure only one-way connections for partial-replicas,
241     # and make sure they point the right way.
242     edge_list = []
243     for edge in output_edges:
244         # We know these edges only have two endpoints because we made
245         # them.
246         v, w = edge.vertices
247         if v.site is my_site or w.site is my_site:
248             if (((v.is_black() or w.is_black()) and
249                  v.dist_to_red != MAX_DWORD)):
250                 edge.directed = True
251
252                 if w.dist_to_red < v.dist_to_red:
253                     edge.vertices[:] = w, v
254             edge_list.append(edge)
255
256     if verify or dot_file_dir is not None:
257         graph_edges = [[x.site.site_dnstr for x in e.vertices]
258                        for e in edge_list]
259         #add the reverse edge if not directed.
260         graph_edges.extend([x.site.site_dnstr
261                             for x in reversed(e.vertices)]
262                            for e in edge_list if not e.directed)
263         graph_nodes = [x.site.site_dnstr for x in graph.vertices]
264         verify_properties = ()
265         verify_and_dot('post-one-way-partial', graph_edges, graph_nodes,
266                        label=label, properties=verify_properties,
267                        debug=DEBUG, verify=verify,
268                        directed=True,
269                        dot_file_dir=dot_file_dir)
270
271     # count the components
272     return edge_list, components
273
274
275 def create_edge(con_type, site_link, guid_to_vertex):
276     e = MultiEdge()
277     e.site_link = site_link
278     e.vertices = []
279     for site_guid in site_link.site_list:
280         if str(site_guid) in guid_to_vertex:
281             e.vertices.extend(guid_to_vertex.get(str(site_guid)))
282     e.repl_info.cost = site_link.cost
283     e.repl_info.options = site_link.options
284     e.repl_info.interval = site_link.interval
285     e.repl_info.schedule = convert_schedule_to_repltimes(site_link.schedule)
286     e.con_type = con_type
287     e.directed = False
288     return e
289
290
291 def create_auto_edge_set(graph, transport):
292     e_set = MultiEdgeSet()
293     # use a NULL guid, not associated with a SiteLinkBridge object
294     e_set.guid = misc.GUID()
295     for site_link in graph.edges:
296         if site_link.con_type == transport:
297             e_set.edges.append(site_link)
298
299     return e_set
300
301
302 def create_edge_set(graph, transport, site_link_bridge):
303     # TODO not implemented - need to store all site link bridges
304     e_set = MultiEdgeSet()
305     # e_set.guid = site_link_bridge
306     return e_set
307
308
309 def setup_vertices(graph):
310     for v in graph.vertices:
311         if v.is_white():
312             v.repl_info.cost = MAX_DWORD
313             v.root = None
314             v.component_id = None
315         else:
316             v.repl_info.cost = 0
317             v.root = v
318             v.component_id = v
319
320         v.repl_info.interval = 0
321         v.repl_info.options = 0xFFFFFFFF
322         v.repl_info.schedule = None  # TODO highly suspicious
323         v.demoted = False
324
325
326 def dijkstra(graph, edge_type, include_black):
327     queue = []
328     setup_dijkstra(graph, edge_type, include_black, queue)
329     while len(queue) > 0:
330         cost, guid, vertex = heapq.heappop(queue)
331         for edge in vertex.edges:
332             for v in edge.vertices:
333                 if v is not vertex:
334                     # add new path from vertex to v
335                     try_new_path(graph, queue, vertex, edge, v)
336
337
338 def setup_dijkstra(graph, edge_type, include_black, queue):
339     setup_vertices(graph)
340     for vertex in graph.vertices:
341         if vertex.is_white():
342             continue
343
344         if (((vertex.is_black() and not include_black)
345              or edge_type not in vertex.accept_black
346              or edge_type not in vertex.accept_red_red)):
347             vertex.repl_info.cost = MAX_DWORD
348             vertex.root = None  # NULL GUID
349             vertex.demoted = True  # Demoted appears not to be used
350         else:
351             heapq.heappush(queue, (vertex.repl_info.cost, vertex.guid, vertex))
352
353
354 def try_new_path(graph, queue, vfrom, edge, vto):
355     newRI = ReplInfo()
356     #This function combines the repl_info and checks is that there is
357     # a valid time frame for which replication can actually occur,
358     # despite being adequately connected
359     intersect = combine_repl_info(vfrom.repl_info, edge.repl_info, newRI)
360
361     # If the new path costs more than the current, then ignore the edge
362     if newRI.cost > vto.repl_info.cost:
363         return
364
365     if newRI.cost < vto.repl_info.cost and not intersect:
366         return
367
368     new_duration = total_schedule(newRI.schedule)
369     old_duration = total_schedule(vto.repl_info.schedule)
370
371     # Cheaper or longer schedule
372     if newRI.cost < vto.repl_info.cost or new_duration > old_duration:
373         vto.root = vfrom.root
374         vto.component_id = vfrom.component_id
375         vto.repl_info = newRI
376         heapq.heappush(queue, (vto.repl_info.cost, vto.guid, vto))
377
378
379 def check_demote_vertex(vertex, edge_type):
380     if vertex.is_white():
381         return
382
383     # Accepts neither red-red nor black edges, demote
384     if ((edge_type not in vertex.accept_black and
385          edge_type not in vertex.accept_red_red)):
386         vertex.repl_info.cost = MAX_DWORD
387         vertex.root = None
388         vertex.demoted = True  # Demoted appears not to be used
389
390
391 def undemote_vertex(vertex):
392     if vertex.is_white():
393         return
394
395     vertex.repl_info.cost = 0
396     vertex.root = vertex
397     vertex.demoted = False
398
399
400 def process_edge_set(graph, e_set, internal_edges):
401     if e_set is None:
402         for edge in graph.edges:
403             for vertex in edge.vertices:
404                 check_demote_vertex(vertex, edge.con_type)
405             process_edge(graph, edge, internal_edges)
406             for vertex in edge.vertices:
407                 undemote_vertex(vertex)
408     else:
409         for edge in e_set.edges:
410             process_edge(graph, edge, internal_edges)
411
412
413 def process_edge(graph, examine, internal_edges):
414     # Find the set of all vertices touches the edge to examine
415     vertices = []
416     for v in examine.vertices:
417         # Append a 4-tuple of color, repl cost, guid and vertex
418         vertices.append((v.color, v.repl_info.cost, v.ndrpacked_guid, v))
419     # Sort by color, lower
420     DEBUG("vertices is %s" % vertices)
421     vertices.sort()
422
423     color, cost, guid, bestv = vertices[0]
424     # Add to internal edges an edge from every colored vertex to bestV
425     for v in examine.vertices:
426         if v.component_id is None or v.root is None:
427             continue
428
429         # Only add edge if valid inter-tree edge - needs a root and
430         # different components
431         if ((bestv.component_id is not None and
432              bestv.root is not None and
433              v.component_id is not None and
434              v.root is not None and
435              bestv.component_id != v.component_id)):
436             add_int_edge(graph, internal_edges, examine, bestv, v)
437
438
439 # Add internal edge, endpoints are roots of the vertices to pass in
440 # and are always red or black
441 def add_int_edge(graph, internal_edges, examine, v1, v2):
442     root1 = v1.root
443     root2 = v2.root
444
445     red_red = False
446     if root1.is_red() and root2.is_red():
447         red_red = True
448
449     if red_red:
450         if ((examine.con_type not in root1.accept_red_red
451              or examine.con_type not in root2.accept_red_red)):
452             return
453     elif (examine.con_type not in root1.accept_black
454           or examine.con_type not in root2.accept_black):
455         return
456
457     ri = ReplInfo()
458     ri2 = ReplInfo()
459
460     # Create the transitive replInfo for the two trees and this edge
461     if not combine_repl_info(v1.repl_info, v2.repl_info, ri):
462         return
463     # ri is now initialized
464     if not combine_repl_info(ri, examine.repl_info, ri2):
465         return
466
467     newIntEdge = InternalEdge(root1, root2, red_red, ri2, examine.con_type,
468                               examine.site_link)
469     # Order by vertex guid
470     #XXX guid comparison using ndr_pack
471     if newIntEdge.v1.ndrpacked_guid > newIntEdge.v2.ndrpacked_guid:
472         newIntEdge.v1 = root2
473         newIntEdge.v2 = root1
474
475     internal_edges.add(newIntEdge)
476
477
478 def kruskal(graph, edges):
479     for v in graph.vertices:
480         v.edges = []
481
482     components = set([x for x in graph.vertices if not x.is_white()])
483     edges = list(edges)
484
485     # Sorted based on internal comparison function of internal edge
486     edges.sort()
487
488     #XXX expected_num_tree_edges is never used
489     expected_num_tree_edges = 0  # TODO this value makes little sense
490
491     count_edges = 0
492     output_edges = []
493     index = 0
494     while index < len(edges):  # TODO and num_components > 1
495         e = edges[index]
496         parent1 = find_component(e.v1)
497         parent2 = find_component(e.v2)
498         if parent1 is not parent2:
499             count_edges += 1
500             add_out_edge(graph, output_edges, e)
501             parent1.component_id = parent2
502             components.discard(parent1)
503
504         index += 1
505
506     return output_edges, len(components)
507
508
509 def find_component(vertex):
510     if vertex.component_id is vertex:
511         return vertex
512
513     current = vertex
514     while current.component_id is not current:
515         current = current.component_id
516
517     root = current
518     current = vertex
519     while current.component_id is not root:
520         n = current.component_id
521         current.component_id = root
522         current = n
523
524     return root
525
526
527 def add_out_edge(graph, output_edges, e):
528     v1 = e.v1
529     v2 = e.v2
530
531     # This multi-edge is a 'real' edge with no GUID
532     ee = MultiEdge()
533     ee.directed = False
534     ee.site_link = e.site_link
535     ee.vertices.append(v1)
536     ee.vertices.append(v2)
537     ee.con_type = e.e_type
538     ee.repl_info = e.repl_info
539     output_edges.append(ee)
540
541     v1.edges.append(ee)
542     v2.edges.append(ee)
543
544
545 def setup_graph(part, site_table, transport_guid, sitelink_table,
546                 bridges_required):
547     """Set up a GRAPH, populated with a VERTEX for each site
548     object, a MULTIEDGE for each siteLink object, and a
549     MUTLIEDGESET for each siteLinkBridge object (or implied
550     siteLinkBridge).
551
552     ::returns: a new graph
553     """
554     guid_to_vertex = {}
555     # Create graph
556     g = IntersiteGraph()
557     # Add vertices
558     for site_guid, site in site_table.items():
559         vertex = Vertex(site, part)
560         vertex.guid = site_guid
561         vertex.ndrpacked_guid = ndr_pack(site.site_guid)
562         g.vertices.add(vertex)
563         guid_vertices = guid_to_vertex.setdefault(site_guid, [])
564         guid_vertices.append(vertex)
565
566     connected_vertices = set()
567
568     for site_link_dn, site_link in sitelink_table.items():
569         new_edge = create_edge(transport_guid, site_link,
570                                guid_to_vertex)
571         connected_vertices.update(new_edge.vertices)
572         g.edges.add(new_edge)
573
574     # If 'Bridge all site links' is enabled and Win2k3 bridges required
575     # is not set
576     # NTDSTRANSPORT_OPT_BRIDGES_REQUIRED 0x00000002
577     # No documentation for this however, ntdsapi.h appears to have:
578     # NTDSSETTINGS_OPT_W2K3_BRIDGES_REQUIRED = 0x00001000
579     if bridges_required:
580         g.edge_set.add(create_auto_edge_set(g, transport_guid))
581     else:
582         # TODO get all site link bridges
583         for site_link_bridge in []:
584             g.edge_set.add(create_edge_set(g, transport_guid,
585                                            site_link_bridge))
586
587     g.connected_vertices = connected_vertices
588
589     return g
590
591
592 class VertexColor(object):
593     (red, black, white, unknown) = range(0, 4)
594
595
596 class Vertex(object):
597     """Class encapsulation of a Site Vertex in the
598     intersite topology replication algorithm
599     """
600     def __init__(self, site, part):
601         self.site = site
602         self.part = part
603         self.color = VertexColor.unknown
604         self.edges = []
605         self.accept_red_red = []
606         self.accept_black = []
607         self.repl_info = ReplInfo()
608         self.root = self
609         self.guid = None
610         self.component_id = self
611         self.demoted = False
612         self.options = 0
613         self.interval = 0
614
615     def color_vertex(self):
616         """Color each vertex to indicate which kind of NC
617         replica it contains
618         """
619         # IF s contains one or more DCs with full replicas of the
620         # NC cr!nCName
621         #    SET v.Color to COLOR.RED
622         # ELSEIF s contains one or more partial replicas of the NC
623         #    SET v.Color to COLOR.BLACK
624         #ELSE
625         #    SET v.Color to COLOR.WHITE
626
627         # set to minimum (no replica)
628         self.color = VertexColor.white
629
630         for dnstr, dsa in self.site.dsa_table.items():
631             rep = dsa.get_current_replica(self.part.nc_dnstr)
632             if rep is None:
633                 continue
634
635             # We have a full replica which is the largest
636             # value so exit
637             if not rep.is_partial():
638                 self.color = VertexColor.red
639                 break
640             else:
641                 self.color = VertexColor.black
642
643     def is_red(self):
644         assert(self.color != VertexColor.unknown)
645         return (self.color == VertexColor.red)
646
647     def is_black(self):
648         assert(self.color != VertexColor.unknown)
649         return (self.color == VertexColor.black)
650
651     def is_white(self):
652         assert(self.color != VertexColor.unknown)
653         return (self.color == VertexColor.white)
654
655
656 class IntersiteGraph(object):
657     """Graph for representing the intersite"""
658     def __init__(self):
659         self.vertices = set()
660         self.edges = set()
661         self.edge_set = set()
662         # All vertices that are endpoints of edges
663         self.connected_vertices = None
664
665
666 class MultiEdgeSet(object):
667     """Defines a multi edge set"""
668     def __init__(self):
669         self.guid = 0  # objectGuid siteLinkBridge
670         self.edges = []
671
672
673 class MultiEdge(object):
674     def __init__(self):
675         self.site_link = None  # object siteLink
676         self.vertices = []
677         self.con_type = None  # interSiteTransport GUID
678         self.repl_info = ReplInfo()
679         self.directed = True
680
681
682 class InternalEdge(object):
683     def __init__(self, v1, v2, redred, repl, eType, site_link):
684         self.v1 = v1
685         self.v2 = v2
686         self.red_red = redred
687         self.repl_info = repl
688         self.e_type = eType
689         self.site_link = site_link
690
691     def __eq__(self, other):
692         return not self < other and not other < self
693
694     def __ne__(self, other):
695         return self < other or other < self
696
697     def __gt__(self, other):
698         return other < self
699
700     def __ge__(self, other):
701         return not self < other
702
703     def __le__(self, other):
704         return not other < self
705
706     # TODO compare options and interval
707     def __lt__(self, other):
708         if self.red_red != other.red_red:
709             return self.red_red
710
711         if self.repl_info.cost != other.repl_info.cost:
712             return self.repl_info.cost < other.repl_info.cost
713
714         self_time = total_schedule(self.repl_info.schedule)
715         other_time = total_schedule(other.repl_info.schedule)
716         if self_time != other_time:
717             return self_time > other_time
718
719         #XXX guid comparison using ndr_pack
720         if self.v1.guid != other.v1.guid:
721             return self.v1.ndrpacked_guid < other.v1.ndrpacked_guid
722
723         if self.v2.guid != other.v2.guid:
724             return self.v2.ndrpacked_guid < other.v2.ndrpacked_guid
725
726         return self.e_type < other.e_type