# This Source Code Form is subject to the terms of the Mozilla Public # License, v. 2.0. If a copy of the MPL was not distributed with this # file, You can obtain one at http://mozilla.org/MPL/2.0/. """ Incremental algorithm for creating a "deterministic acyclic finite state automaton" (DAFSA). At the time of writing this algorithm, there was existing logic that depended on a different format for the DAFSA, so this contains convenience functions for converting to a compatible structure. This legacy format is defined in make_dafsa.py. """ from typing import Callable, Dict, List, Optional class Node: children: Dict[str, "Node"] parents: Dict[str, List["Node"]] character: str is_root_node: bool is_end_node: bool def __init__(self, character, is_root_node=False, is_end_node=False): self.children = {} self.parents = {} self.character = character self.is_root_node = is_root_node self.is_end_node = is_end_node def __str__(self): """Produce a helpful string representation of this node. This is expected to only be used for debugging. The produced output is: "c[def.] <123>" ^ ^ ^ | | Internal python ID of the node (used for de-duping) | | | One possible path through the tree to the end | Current node character """ if self.is_root_node: return "" elif self.is_end_node: return "" first_potential_match = "" node = self while node.children: first_character = next(iter(node.children)) if first_character: first_potential_match += first_character node = node.children[first_character] return "%s[%s] <%d>" % (self.character, first_potential_match, id(self)) def add_child(self, child): self.children[child.character] = child child.parents.setdefault(self.character, []) child.parents[self.character].append(self) def remove(self): # remove() must only be called when this node has only a single parent, and that # parent doesn't need this child anymore. # The caller is expected to have performed this validation. # (placing asserts here add a non-trivial performance hit) # There's only a single parent, so only one list should be in the "parents" map parent_list = next(iter(self.parents.values())) self.remove_parent(parent_list[0]) for child in list(self.children.values()): child.remove_parent(self) def remove_parent(self, parent_node: "Node"): parent_node.children.pop(self.character) parents_for_character = self.parents[parent_node.character] parents_for_character.remove(parent_node) if not parents_for_character: self.parents.pop(parent_node.character) def copy_fork_node(self, fork_node: "Node", child_to_avoid: Optional["Node"]): """Shallow-copy a node's children. When adding a new word, sometimes previously-joined suffixes aren't perfect matches any more. When this happens, some nodes need to be "copied" out. For all non-end nodes, there's a child to avoid in the shallow-copy. """ for child in fork_node.children.values(): if child is not child_to_avoid: self.add_child(child) def is_fork(self): """Check if this node has multiple parents""" if len(self.parents) == 0: return False if len(self.parents) > 1: return True return len(next(iter(self.parents.values()))) > 1 def is_replacement_for_prefix_end_node(self, old: "Node"): """Check if this node is a valid replacement for an old end node. A node is a valid replacement if it maintains all existing child paths while adding the new child path needed for the new word. Args: old: node being replaced Returns: True if this node is a valid replacement node. """ if len(self.children) != len(old.children) + 1: return False for character, other_node in old.children.items(): this_node = self.children.get(character) if other_node is not this_node: return False return True def is_replacement_for_prefix_node(self, old: "Node"): """Check if this node is a valid replacement for a non-end node. A node is a valid replacement if it: * Has one new child that the old node doesn't * Is missing a child that the old node has * Shares all other children Returns: True if this node is a valid replacement node. """ if len(self.children) != len(old.children): return False found_extra_child = False for character, other_node in old.children.items(): this_node = self.children.get(character) if other_node is not this_node: if found_extra_child: # Found two children in the old node that aren't in the new one, # this isn't a valid replacement return False else: found_extra_child = True return found_extra_child class SuffixCursor: index: int # Current position of the cursor within the DAFSA. node: Node def __init__(self, index, node): self.index = index self.node = node def _query(self, character: str, check: Callable[[Node], bool]): for node in self.node.parents.get(character, []): if check(node): self.index -= 1 self.node = node return True return False def find_single_child(self, character): """Find the next matching suffix node that has a single child. Return True if such a node is found.""" return self._query(character, lambda node: len(node.children) == 1) def find_end_of_prefix_replacement(self, end_of_prefix: Node): """Find the next matching suffix node that replaces the old prefix-end node. Return True if such a node is found.""" return self._query( end_of_prefix.character, lambda node: node.is_replacement_for_prefix_end_node(end_of_prefix), ) def find_inside_of_prefix_replacement(self, prefix_node: Node): """Find the next matching suffix node that replaces a node within the prefix. Return True if such a node is found.""" return self._query( prefix_node.character, lambda node: node.is_replacement_for_prefix_node(prefix_node), ) class DafsaAppendStateMachine: """State machine for adding a word to a Dafsa. Each state returns a function reference to the "next state". States should be invoked until "None" is returned, in which case the new word has been appended. The prefix and suffix indexes are placed according to the currently-known valid value (not the next value being investigated). Additionally, they are 0-indexed against the root node (which sits behind the beginning of the string). Let's imagine we're at the following state when adding, for example, the word "mozilla.org": mozilla.org ^ ^ ^ ^ | | | | / | | \ [root] | | [end] node node | \ | suffix \ prefix In this state, the furthest prefix match we could find was: [root] - m - o - z - i - l The index of the prefix match is "5". Additionally, we've been looking for suffix nodes, and we've already found: r - g - [end] The current suffix index is "10". The next suffix node we'll attempt to find is at index "9". """ root_node: Node prefix_index: int suffix_cursor: SuffixCursor stack: List[Node] word: str suffix_overlaps_prefix: bool first_fork_index: Optional[int] _state: Callable def __init__(self, word, root_node, end_node): self.root_node = root_node self.prefix_index = 0 self.suffix_cursor = SuffixCursor(len(word) + 1, end_node) self.stack = [root_node] self.word = word self.suffix_overlaps_prefix = False self.first_fork_index = None self._state = self._find_prefix def run(self): """Run this state machine to completion, adding the new word.""" while self._state is not None: self._state = self._state() def _find_prefix(self): """Find the longest existing prefix that matches the current word.""" prefix_node = self.root_node while self.prefix_index < len(self.word): next_character = self.word[self.prefix_index] next_node = prefix_node.children.get(next_character) if not next_node: # We're finished finding the prefix, let's find the longest suffix # match now. return self._find_suffix_nodes_after_prefix self.prefix_index += 1 prefix_node = next_node self.stack.append(next_node) if not self.first_fork_index and next_node.is_fork(): self.first_fork_index = self.prefix_index # Deja vu, we've appended this string before. Since this string has # already been appended, we don't have to do anything. return None def _find_suffix_nodes_after_prefix(self): """Find the chain of suffix nodes for characters after the prefix.""" while self.suffix_cursor.index - 1 > self.prefix_index: # To fetch the next character, we need to subtract two from the current # suffix index. This is because: # * The next suffix node is 1 node before our current node (subtract 1) # * The suffix index includes the root node before the beginning of the # string - it's like the string is 1-indexed (subtract 1 again). next_character = self.word[self.suffix_cursor.index - 2] if not self.suffix_cursor.find_single_child(next_character): return self._add_new_nodes if self.suffix_cursor.node is self.stack[-1]: # The suffix match is overlapping with the prefix! This can happen in # cases like: # * "ab" # * "abb" # The suffix cursor is at the same node as the prefix match, but they're # at different positions in the word. # # [root] - a - b - [end] # ^ # / \ # / \ # prefix suffix # \ / # \ / # VV # "abb" if not self.first_fork_index: # There hasn't been a fork, so our prefix isn't shared. So, we # can mark this node as a fork, since the repetition means # that there's two paths that are now using this node self.first_fork_index = self.prefix_index return self._add_new_nodes # Removes the link between the unique part of the prefix and the # shared part of the prefix. self.stack[self.first_fork_index].remove_parent( self.stack[self.first_fork_index - 1] ) self.suffix_overlaps_prefix = True if self.first_fork_index is None: return self._find_next_suffix_nodes elif self.suffix_cursor.index - 1 == self.first_fork_index: return self._find_next_suffix_node_at_prefix_end_at_fork else: return self._find_next_suffix_node_at_prefix_end_after_fork def _find_next_suffix_node_at_prefix_end_at_fork(self): """Find the next suffix node that replaces the end of the prefix. In this state, the prefix_end node is the same as the first fork node. Therefore, if a match can be found, the old prefix node can't be entirely deleted since it's used elsewhere. Instead, just the link between our unique prefix and the end of the fork is removed. """ existing_node = self.stack[self.prefix_index] if not self.suffix_cursor.find_end_of_prefix_replacement(existing_node): return self._add_new_nodes self.prefix_index -= 1 self.first_fork_index = None if not self.suffix_overlaps_prefix: existing_node.remove_parent(self.stack[self.prefix_index]) else: # When the suffix overlaps the prefix, the old "parent link" was removed # earlier in the "find_suffix_nodes_after_prefix" step. self.suffix_overlaps_prefix = False return self._find_next_suffix_nodes def _find_next_suffix_node_at_prefix_end_after_fork(self): """Find the next suffix node that replaces the end of the prefix. In this state, the prefix_end node is after the first fork node. Therefore, even if a match is found, we don't want to modify the replaced prefix node since an unrelated word chain uses it. """ existing_node = self.stack[self.prefix_index] if not self.suffix_cursor.find_end_of_prefix_replacement(existing_node): return self._add_new_nodes self.prefix_index -= 1 if self.prefix_index == self.first_fork_index: return self._find_next_suffix_node_within_prefix_at_fork else: return self._find_next_suffix_nodes_within_prefix_after_fork def _find_next_suffix_node_within_prefix_at_fork(self): """Find the next suffix node within a prefix. In this state, we've already worked our way back and found nodes in the suffix to replace prefix nodes after the fork node. We have now reached the fork node, and if we find a replacement for it, then we can remove the link between it and our then-unique prefix and clear the fork status. """ existing_node = self.stack[self.prefix_index] if not self.suffix_cursor.find_inside_of_prefix_replacement(existing_node): return self._add_new_nodes self.prefix_index -= 1 self.first_fork_index = None if not self.suffix_overlaps_prefix: existing_node.remove_parent(self.stack[self.prefix_index]) else: # When the suffix overlaps the prefix, the old "parent link" was removed # earlier in the "find_suffix_nodes_after_prefix" step. self.suffix_overlaps_prefix = False return self._find_next_suffix_nodes def _find_next_suffix_nodes_within_prefix_after_fork(self): """Find the next suffix nodes within a prefix. Finds suffix nodes to replace prefix nodes, but doesn't modify the prefix nodes since they're after a fork (so, we're sharing prefix nodes with other words and can't modify them). """ while True: existing_node = self.stack[self.prefix_index] if not self.suffix_cursor.find_inside_of_prefix_replacement(existing_node): return self._add_new_nodes self.prefix_index -= 1 if self.prefix_index == self.first_fork_index: return self._find_next_suffix_node_within_prefix_at_fork def _find_next_suffix_nodes(self): """Find all remaining suffix nodes in the chain. In this state, there's no (longer) any fork, so there's no other words using our current prefix. Therefore, as we find replacement nodes as we work our way backwards, we can remove the now-unused prefix nodes. """ while True: existing_node = self.stack[self.prefix_index] if not self.suffix_cursor.find_end_of_prefix_replacement(existing_node): return self._add_new_nodes # This prefix node is wholly replaced by the new suffix node, so it can # be deleted. existing_node.remove() self.prefix_index -= 1 def _add_new_nodes(self): """Adds new nodes to support the new word. Duplicates forked nodes to make room for new links, adds new nodes for new characters, and splices the prefix to the suffix to finish embedding the new word into the DAFSA. """ if self.first_fork_index is not None: front_node = _duplicate_fork_nodes( self.stack, self.first_fork_index, self.prefix_index, # if suffix_overlaps_parent, the parent link was removed # earlier in the word-adding process. remove_parent_link=not self.suffix_overlaps_prefix, ) else: front_node = self.stack[self.prefix_index] new_text = self.word[self.prefix_index : self.suffix_cursor.index - 1] for character in new_text: new_node = Node(character) front_node.add_child(new_node) front_node = new_node front_node.add_child(self.suffix_cursor.node) return None # Done! def _duplicate_fork_nodes(stack, fork_index, prefix_index, remove_parent_link=True): parent_node = stack[fork_index - 1] if remove_parent_link: # remove link to old chain that we're going to be copying stack[fork_index].remove_parent(parent_node) for index in range(fork_index, prefix_index + 1): fork_node = stack[index] replacement_node = Node(fork_node.character) child_to_avoid = None if index < len(stack) - 1: # We're going to be manually replacing the next node in the stack, # so don't connect it as a child. child_to_avoid = stack[index + 1] replacement_node.copy_fork_node(fork_node, child_to_avoid) parent_node.add_child(replacement_node) parent_node = replacement_node return parent_node class Dafsa: root_node: Node end_node: Node def __init__(self): self.root_node = Node(None, is_root_node=True) self.end_node = Node(None, is_end_node=True) @classmethod def from_tld_data(cls, lines): """Create a dafsa for TLD data. TLD data has a domain and a "type" enum. The source data encodes the type as a text number, but the dafsa-consuming code assumes that the type is a raw byte number (e.g.: "1" => 0x01). This function acts as a helper, processing this TLD detail before creating a standard dafsa. """ dafsa = cls() for i, word in enumerate(lines): domain_number = word[-1] # Convert type from string to byte representation raw_domain_number = chr(ord(domain_number) & 0x0F) word = "%s%s" % (word[:-1], raw_domain_number) dafsa.append(word) return dafsa def append(self, word): state_machine = DafsaAppendStateMachine(word, self.root_node, self.end_node) state_machine.run()