umnetdb-utils 0.1.5__py3-none-any.whl → 0.2.0__py3-none-any.whl

umnetdb_utils/cli.py

@@ -0,0 +1,114 @@
+import typer
+from umnetdb_utils import UMnetdb
+import inspect
+from functools import wraps
+
+from rich.console import Console
+from rich.table import Table
+
+from typing import Callable, List
+from typing_extensions import Annotated
+
+app = typer.Typer()
+
+def print_result(result:List[dict]):
+    """
+    Takes the result of a umnetdb call and prints it as a table
+    """
+    if len(result) == 0:
+        print("No results found")
+        return
+    
+    if isinstance(result, dict):
+        result = [result]
+        
+    # instantiate table with columns based on entry dict keys
+    table = Table(*result[0].keys())
+    for row in result:
+        table.add_row(*[str(i) for i in row.values()])
+
+    console = Console()
+    console.print(table)
+    
+
+def command_generator(method_name:str, method:Callable):
+    """
+    Generates a typer command function for an arbitrary method
+    in the umnetdb class. The generated function opens a connection with 
+    the database, executes the method, and prints out the results.
+
+    Note that the docstring of each method is interrogated to generate
+    help text for each typer command.
+
+    :method_name: The name of the method
+    :method: The method itself
+    """
+
+    # first we're going to tease out the 'help' portions of the method
+    # from the docstring.
+    docstr = method.__doc__
+    docstr_parts = docstr.split("\n:")
+
+    # first section of the docstring is always a generic 'this is what the method does'.
+    cmd_help = docstr_parts.pop(0)
+
+    # next sections are details on the specific arguments that we want to pass to typer as
+    # special annotated type hints
+    arg_help = {}
+    for arg_str in docstr_parts:
+        if ":" in arg_str:
+            arg, help = arg_str.split(":")
+            arg_help[arg] = help.strip()
+        
+    sig = inspect.signature(method)
+
+    # going through the method's arguments and augmenting the 'help' section for each one
+    # from the docstring if applicable
+    new_params = []
+    for p_name, p in sig.parameters.items():
+
+        # need to skip self
+        if p_name == "self":
+            continue
+
+        # if there wasn't any helper text then just append the parameter as is
+        if p_name not in arg_help:
+            new_params.append(p)
+            continue
+
+        # params without default values should be typer 'arguments'
+        if p.default == inspect._empty:
+            new_params.append(p.replace(annotation=Annotated[p.annotation, typer.Argument(help=arg_help[p_name])]))
+            continue
+
+        # params with default values should be typer 'options'
+        new_params.append(p.replace(annotation=Annotated[p.annotation, typer.Option(help=arg_help[p_name])]))
+
+    new_sig = sig.replace(parameters=new_params)
+
+
+    # new munged function based on the origional method, with a new signature
+    # and docstring for typer
+    @wraps(method)
+    def wrapper(*args, **kwargs):
+        with UMnetdb() as db:
+            result = getattr(db, method_name)(*args, **kwargs)
+        print_result(result)
+    
+    wrapper.__signature__ = new_sig
+    wrapper.__doc__ = cmd_help
+
+    return wrapper
+
+
+def main():
+    for f_name,f in UMnetdb.__dict__.items():
+        if not(f_name.startswith("_")) and callable(f):
+            app.command()(command_generator(f_name, f))
+
+    app()
+
+if __name__ == "__main__":
+    main()
+
+    

umnetdb_utils/umnetdb.py

@@ -1,16 +1,18 @@
-from typing import List
+from typing import List, Optional
 import logging
 import re
+import ipaddress
+from copy import deepcopy
 
-
-from sqlalchemy import text
 from .base import UMnetdbBase
+from .utils import is_ip_address, Packet, Hop, Path, LOCAL_PROTOCOLS
 
+logger = logging.getLogger(__name__)
 
 class UMnetdb(UMnetdbBase):
     URL = "postgresql+psycopg://{UMNETDB_USER}:{UMNETDB_PASSWORD}@wintermute.umnet.umich.edu/umnetdb"
 
-    def get_neighbors(self, device: str, known_devices_only: bool = True) -> List[dict]:
+    def get_neighbors(self, device: str, known_devices_only: bool = True, interface:Optional[str]=None) -> List[dict]:
         """
         Gets a list of the neighbors of a particular device. If the port
         has a parent in the LAG table that is included as well.
@@ -20,9 +22,12 @@ class UMnetdb(UMnetdbBase):
 
         Setting 'known_devices_only' to true only returns neighbors that are found
         in umnet_db's device table. Setting it to false will return all lldp neighbors
-        and will include things like phones and APs
+        and will include things like phones and APs.
 
         Returns results as a list of dictionary entries keyed on column names.
+        :device: Name of the device
+        :known_devices_only: If set to true, will only return neighbors found in umnetdb's device table.
+        :interface: If supplied, restrict to only find neighbors on a particular interface on the device
         """
 
         if known_devices_only:
@@ -55,6 +60,8 @@ class UMnetdb(UMnetdbBase):
         table = "neighbor n"
 
         where = [f"n.device='{device}'"]
+        if interface:
+            where.append(f"n.port='{interface}'")
 
         query = self._build_select(select, table, joins, where)
 
@@ -64,9 +71,11 @@ class UMnetdb(UMnetdbBase):
         """
         Gets all devices within a DL zone based on walking the 'neighbors'
         table.
-
+        
         For each device, the following attributes are returned:
         "name", "ip", "version", "vendor", "model", "serial"
+
+        :zone_name: Name of the DL zone
         """
         device_cols = ["name", "ip", "version", "vendor", "model", "serial"]
 
@@ -90,13 +99,12 @@ class UMnetdb(UMnetdbBase):
         while len(todo) != 0:
             device = todo.pop()
 
-            print(f"Processing device {device}")
             # note that by default this method only returns neighbors in the 'device' table,
             # any others are ignored
             neighs = self.get_neighbors(device)
             devices_by_name[device]["neighbors"] = {}
             for neigh in neighs:
-                print(f"EVALUATING {neigh}")
+
                 # only want 'd- or 'dl-' or 's-' devices, and we don't want out of band devices
                 if re.match(r"(dl?-|s-)", neigh["remote_device"]) and not re.match(
                     r"s-oob-", neigh["remote_device"]
@@ -120,3 +128,291 @@ class UMnetdb(UMnetdbBase):
                         todo.append(neigh_device["name"])
 
         return list(devices_by_name.values())
+
+
+    def l3info(self, search_str:str, detail:bool=False, num_results:int=10, exact:bool=False)->list[dict]:
+        """
+        Does a search of the umnetdb ip_interface table.
+        
+        :search_str: Can be an IP address, 'VlanX', or a full or partial netname
+        :detail: Adds admin/oper status, primary/secondary, helpers, and timestamps to output.
+        :num_results: Limits number of results printed out.
+        :exact: Only return exact matches, either for IP addresses or for string matches.
+        """
+
+        cols = ["device", "ip_address", "interface", "description", "vrf"]
+        if detail:
+            cols.extend(["admin_up", "oper_up", "secondary", "helpers", "first_seen", "last_updated"])
+
+        # 'is contained within' IP search - reference:
+        # https://www.postgresql.org/docs/9.3/functions-net.html
+
+        # VlanX based searches are always 'exact'
+        if re.match(r"Vlan\d+$", search_str):
+                where = [f"interface = '{search_str}'"]
+
+        # ip or description based searches can be 'exact' or inexact
+        elif exact:
+            if is_ip_address(search_str):
+                where = [f"host(ip_address) = '{search_str}'"]
+            else:
+                where = [f"description = '{search_str}'"]
+
+        else:
+            if is_ip_address(search_str):
+                where = [f"ip_address >>= '{search_str}'"]
+            else:
+                where = [f"description like '{search_str}%'"]
+        
+        # removing IPs assigned to mgmt interfaces
+        where.append("vrf != 'management'")
+
+        query = self._build_select(
+            select=cols,
+            table="ip_interface",
+            where=where,
+            limit=num_results,
+        )
+        return self.execute(query)
+
+    
+    def route(self, router:str, prefix:str, vrf:str, resolve_nh:bool=True, details:bool=False) -> list[dict]:
+        """
+        Does an lpm query on a particular router for a particular prefix
+        in a particular VRF.
+
+        :router: Name of the router to query
+        :prefix: Prefix to query for
+        :vrf: Name of the VRF to query against
+        :resolve_nh: If no nh_interface is present in the database, recursively resolve for it.
+        :details: Set to true to get output of all columns in the route table.
+        """
+        if details:
+            cols = ["*"]
+        else:
+            cols = ["device", "vrf", "prefix", "nh_ip", "nh_table", "nh_interface"]
+            
+        lpms = self.lpm_query(router, prefix, vrf, columns=cols)
+
+        if not lpms:
+            return []
+        
+        if not resolve_nh:
+            return lpms
+
+        resolved_results = []
+        for route in lpms:
+            if route["nh_interface"]:
+                resolved_results.append(route)
+            else:
+                self._resolve_nh(route, resolved_results, 0)
+
+        return resolved_results
+
+    def lpm_query(self, router:str, prefix:str, vrf:str, columns:Optional[str]=None)->list[dict]:
+        """
+        Does an lpm query against a particular router, prefix, and vrf. Optionally specify
+        which columns you want to limit the query to.
+        """
+
+        select = columns if columns else ["*"]
+
+        query = self._build_select(
+            select=select,
+            table="route",
+            where=[f"device='{router}'", f"prefix >>= '{prefix}'", f"vrf='{vrf}'"],
+            order_by="prefix"
+        )
+
+        result = self.execute(query)
+
+        if not result:
+            return None
+        
+        if len(result) == 1:
+            return result
+        
+        # peeling the longest matching prefix of the end of the results, which
+        # are ordered by ascending prefixlength
+        lpm_results = [result.pop()]
+
+        # finding any other equivalent lpms. As soon as we run into one that
+        # doesn't match we know we're done.
+        result.reverse()
+        for r in result:
+            if r["prefix"] == lpm_results[0]["prefix"]:
+                lpm_results.append(r)
+            else:
+                break
+    
+        return lpm_results
+    
+    def mpls_label(self, router:str, label:str) -> list[dict]:
+        """
+        Looks up a particular label for a particular device in the mpls table.
+        :router: device name
+        :label: label value
+        """
+        query = self._build_select(
+            select=["*"],
+            table="mpls",
+            where=[f"device='{router}'", f"in_label='{label}'"]
+        )
+        return self.execute(query)
+    
+    def vni(self, router:str, vni:int) -> dict:
+        """
+        Looks up a particular vni on the router and returns
+        the VRF or vlan_id it's associated with
+        """
+
+        query = self._build_select(
+            select=["*"],
+            table="vni",
+            where=[f"device='{router}'", f"vni='{vni}'"]
+        )
+        return self.execute(query, fetch_one=True)
+
+    
+    def _resolve_nh(self, route:dict, resolved_routes:list[dict], depth:int) -> dict:
+        """
+        Recursively resolves next hop of a route till we find a nh interface.
+        If we hit a recursion depth of 4 then an exception is thrown - the max depth on our
+        network I've seen is like 2 (for a static route to an indirect next hop)
+        """
+        if route["nh_interface"]:
+            return route
+        
+        depth += 1        
+        if depth == 4:
+            raise RecursionError(f"Reached max recursion depth of 4 trying to resolve route {route}")
+        
+        nh_ip = route["nh_ip"]
+        nh_table = route["nh_table"]
+        router = route["device"]
+
+        nh_routes = self.lpm_query(router, nh_ip, nh_table)
+
+        for nh_route in nh_routes:
+            r = self._resolve_nh(nh_route, resolved_routes, depth)
+            resolved_route = deepcopy(route)
+            resolved_route["nh_interface"] = r["nh_interface"]
+            if r["mpls_label"] and resolved_route["mpls_label"] is not None:
+                resolved_route["mpls_label"].extend(r["mpls_label"])
+            elif r["mpls_label"]:
+                resolved_route["mpls_label"] = r["mpls_label"]
+            resolved_route["nh_ip"] = r["nh_ip"]
+            resolved_routes.append(resolved_route)
+
+
+    def get_all_paths(self, src_ip:str, dst_ip:str):
+        """
+        Traces the path between a particular source and destination IP
+        :src_ip: A source IP address, must be somewhere on our network
+        :dst_ip: A destination IP address, does not have to be on our network.
+        """
+        for ip_name, ip in [("source", src_ip), ("destination", dst_ip)]:
+            if not is_ip_address(ip):
+                raise ValueError(f"invalid {ip_name} IP address {ip}")
+        
+        src_l3info = self.l3info(src_ip, num_results=1)
+        if not src_l3info:
+            raise ValueError(f"Could not find where {src_ip} is routed")
+        
+        packet = Packet(dst_ip=ipaddress.ip_address(dst_ip))
+        hop = Hop(src_l3info[0]["device"], src_l3info[0]["vrf"], src_l3info[0]["interface"], packet)
+        path = Path(hop)
+
+        self._walk_path(path, hop, hop.router, hop.vrf, packet)
+
+        return path.get_path()
+
+    
+    def _walk_path(self, path:Path, curr_hop:Hop, nh_router:str, nh_table:str, packet:Packet):
+
+        logger.debug(f"\n******* walking path - current hop: {curr_hop}, nh_route: {nh_router}, nh_table {nh_table} *******")
+        logger.debug(f"Known hops: {path.hops.keys()}")
+
+        # mpls-based lookup
+        if packet.label_stack:
+            logger.debug("")
+            routes = self.mpls_label(router=nh_router, label=packet.label_stack[-1])
+
+        # otherwise we want to do an ip based lokup
+        else:
+            routes = self.route(router=nh_router, prefix=packet.dst_ip, vrf=nh_table, details=True)
+
+        if not routes:
+            raise ValueError(f"No route found for {curr_hop}")
+
+        for idx, route in zip(range(1,len(routes)+1), routes):
+
+            logger.debug(f"*** Processing route {idx} of {len(routes)}:{route} ***")
+            nh_router = None
+            nh_table = None
+            new_packet = deepcopy(packet)
+
+            # if the packet is not encapsulated and the route is local, we have reached our destination
+            if not packet.is_encapped() and route.get("protocol") in LOCAL_PROTOCOLS:
+                logger.debug(f"Destination reached at {curr_hop}")
+                final_hop = Hop(route["device"], vrf=route["vrf"], interface=route["nh_interface"], packet=new_packet)
+                path.add_hop(curr_hop, final_hop)
+                continue
+
+            # VXLAN decap - requires local lookup in the vrf that maps to the packet's VNI
+            if route.get("protocol") == "direct" and packet.vni and curr_hop.router == route["device"]:
+                vni = self.vni(router=route["device"], vni=packet.vni)
+                new_packet.vxlan_decap()
+                nh_table = vni["vrf"]
+                nh_router = route["device"]
+                logger.debug(f"vxlan-decapping packet, new packet {new_packet}")
+
+            # MPLS decap - requires local lookup in the vrf indicated by 'nh_interface' field
+            # of this aggregate route
+            elif route.get("aggregate"):
+                new_packet.mpls_pop()
+                nh_table = route["nh_interface"]
+                nh_router = route["device"]
+                logger.debug(f"mpls aggregate route, new packet {new_packet}")
+
+            # VXLAN encap - requires local lookup of encapped packet in the nh table
+            elif route.get("vxlan_vni") and not packet.is_encapped():
+                new_packet.vxlan_encap(route["vxlan_vni"], route["vxlan_endpoint"])
+                logger.debug(f"vxlan-encapping packet, new packet {new_packet}")
+
+            # MPLS encap for an IP route. Resolved routes will have both transport and vrf
+            # labels if applicable - this 'push' will add both to the packet.
+            elif route.get("mpls_label"):
+                new_packet.mpls_push(route["mpls_label"])
+                logger.debug(f"mpls-encapping packet, new packet {new_packet}")
+
+            # MPLS route - note in our environment we don't have anything that requires a
+            # push on an already-labeled packet (!)
+            elif route.get("in_label"):
+
+                if route["out_label"] == ["pop"]:
+                    new_packet.mpls_pop()
+                else:
+                    new_packet.mpls_swap(route["out_label"])
+                logger.debug(f"mpls push or swap, new packet {new_packet}")
+            
+            # if the next hop isn't local we need to figure out which router it's on. In our environment
+            # the easiest way to do that is to use l3info against the nh_ip of the route.
+            if not nh_router and route["nh_ip"]:
+                
+                l3i_router = self.l3info(str(route["nh_ip"]), exact=True)
+                if l3i_router:
+                    nh_router = l3i_router[0]["device"]
+                    nh_table = nh_table if nh_table else l3i_router[0]["vrf"]
+                    logger.debug(f"found router {nh_router} for nh ip {route['nh_ip']}")
+
+            if not nh_router:
+                raise ValueError(f"Unknown next hop for {curr_hop} route {route}")
+
+            # add this hop to our path and if it's a new hop, keep waking
+            new_hop = Hop(route["device"], vrf=route.get("nh_table", "default"), interface=route["nh_interface"], packet=new_packet)
+            logger.debug(f"new hop generated: {new_hop}")
+            new_path = path.add_hop(curr_hop, new_hop)
+            if new_path:
+                logger.debug("New path detected - still walking")
+                self._walk_path(path, new_hop, nh_router, nh_table, new_packet)

umnetdb_utils/utils.py

@@ -1,8 +1,26 @@
 import ipaddress
 import re
 
+from dataclasses import dataclass, field
+from copy import copy, deepcopy
 
-def is_ip_address(input_str, version=None):
+from typing import Optional, Union, List, Dict
+
+type IPAddress = Union[ipaddress.IPv4Address, ipaddress.IPv6Address]
+
+LOCAL_PROTOCOLS = [
+    "Access-internal",
+    "am",
+    "connected",
+    "direct",
+    "Direct",
+    "hmm",
+    "local",
+    "VPN",
+    "vrrpv3",
+]
+
+def is_ip_address(input_str:str, version:Optional[int]=None):
     try:
         ip = ipaddress.ip_address(input_str)
     except ValueError:
@@ -14,7 +32,7 @@ def is_ip_address(input_str, version=None):
     return True
 
 
-def is_ip_network(input_str, version=None):
+def is_ip_network(input_str:str, version:Optional[int]=None):
     # First check that this is a valid IP or network
     try:
         net = ipaddress.ip_network(input_str)
@@ -27,7 +45,7 @@ def is_ip_network(input_str, version=None):
     return True
 
 
-def is_mac_address(input_str):
+def is_mac_address(input_str:str):
     """
     Validates the input string as a mac address. Valid formats are
     XX:XX:XX:XX:XX:XX, XX-XX-XX-XX-XX-XX, XXXX.XXXX.XXXX
@@ -40,3 +58,266 @@ def is_mac_address(input_str):
         return True
 
     return False
+
+
+@dataclass(order=True)
+class Packet:
+    """
+    Modeling a packet as it traverses the network
+    """
+
+    dst_ip: IPAddress
+    #ttl: int = 255
+    inner_dst_ip: Optional[IPAddress] = None
+    label_stack: Optional[list[int]] = None
+    vni: Optional[int] = None
+
+    def __str__(self):
+        if self.label_stack:
+            return f"MPLS: {self.label_stack}"
+        if self.vni:
+            return f"VXLAN: {self.dst_ip}:{self.vni}"
+        return f"{self.dst_ip}"
+
+    # def decrement_ttl(self):
+    #     new_packet = copy(self)
+    #     new_packet.ttl -= 1
+    #     return new_packet
+
+    def vxlan_encap(self, vni:int, tunnel_destination: ipaddress.IPv4Address) -> "Packet":
+        """
+        Return a copy of the existing packet, but with a vxlan encap
+        """
+        if self.vni or self.inner_dst_ip or self.label_stack:
+            raise ValueError(f"Can't encapsulate an already-encapsulated packet: {self}")
+        
+        # new_packet = deepcopy(self)
+        # new_packet.inner_dst_ip = self.dst_ip
+        # new_packet.vni = vni
+        # new_packet.dst_ip = tunnel_destination
+
+        # return new_packet
+        self.inner_dst_ip = self.dst_ip
+        self.vni = vni
+        self.dst_ip = tunnel_destination
+    
+    def vxlan_decap(self) -> "Packet":
+        """
+        Return a copy of the existing packet, but with a vxlan decap
+        """
+        if not self.vni or not self.inner_dst_ip or self.label_stack:
+            raise ValueError(f"Can't decap a non-VXLAN packet: {self}")
+        
+        # new_packet = deepcopy(self)
+        # new_packet.dst_ip = self.inner_dst_ip
+        # new_packet.vni = None
+        # new_packet.inner_dst_ip = None
+
+        # return new_packet
+        self.dst_ip = self.inner_dst_ip
+        self.vni = None
+        self.inner_dst_ip = None
+
+    def mpls_push(self, label:Union[int,list[int]]) -> "Packet":
+        """
+        Retrun a copy of the existing packet but with MPLS labels pushed on the stack
+        """
+        if self.vni or self.inner_dst_ip:
+            raise ValueError(f"Can't do mpls on a VXLAN packet: {self}")
+        
+        # new_packet = deepcopy(self)
+        # if not new_packet.label_stack:
+        #     new_packet.label_stack = []
+        
+        # if isinstance(label, int):
+        #     new_packet.label_stack.append(label)
+        # else:
+        #     new_packet.label_stack.extend(label)
+
+        # return new_packet
+        if not self.label_stack:
+            self.label_stack = []
+        
+        if isinstance(label, list):
+            self.label_stack.extend(label)
+        else:
+            self.label_stack.append(label)
+
+    def mpls_pop(self, num_pops:int=1) -> "Packet":
+        """
+        Return a copy of the existing packet but with MPLS label(s) popped
+        """
+        if not self.label_stack:
+            raise ValueError(f"Can't pop from an empty label stack!: {self}")
+        if len(self.label_stack) < num_pops:
+            raise ValueError(f"Can't pop {num_pops} labels from packet: {self}")
+        
+        # new_packet = copy(self)
+        # for _ in range(num_pops):
+        #     new_packet.label_stack.pop()
+
+        # return new_packet
+
+        for _ in range(num_pops):
+            self.label_stack.pop()
+
+
+    def mpls_swap(self, label:Union[int,list[int]]) -> "Packet":
+        """
+        Rerturn a copy of the existing packet but with MPLS label swap
+        """
+        if not self.label_stack:
+            raise ValueError(f"Can't pop from an empty label stack!: {self}")   
+
+        # new_packet = copy(self)
+        # new_packet.label_stack.pop()
+        # if isinstance(label, int):
+        #     new_packet.label_stack.append(label)
+
+        # # inconsistency in mpls table where this really should be a 'pop' but
+        # # it shows up as a single blank entry in a list.
+        # elif label == [""]:
+        #     new_packet.label_stack = []
+        # else:
+        #     new_packet.label_stack.extend(label)   
+        
+        # return new_packet
+        self.label_stack.pop()
+        if label == [""]:
+            self.lable_stack = []
+        elif isinstance(label, list):
+            self.label_stack.extend(label)
+        else:
+            self.label_stack.append(label)
+    
+    def is_encapped(self):
+        return bool(self.vni or self.label_stack)
+        
+
+@dataclass(order=True)
+class Hop:
+    """
+    A hop along the path and the corresponding packet at that location
+    """
+    router: str
+    vrf: str
+    interface: str
+    packet: Packet
+
+    _parent: "Hop" = field(init=False, compare=False, default=None)
+    _children: List["Hop"] = field(init=False, compare=False, default_factory=list)
+
+    def __str__(self):
+        vrf = "" if self.vrf == "default" else f" VRF {self.vrf}"
+        return f"{self.router}{vrf} {self.interface}, Packet: {self.packet}"
+
+    def copy(self):
+        return deepcopy(self)
+    
+    @classmethod
+    def unknown(self, packet=Packet):
+        """
+        Returns an "unknown" hop for a specific packet
+        """
+        return Hop("","","", packet)
+    
+    def is_unknown(self):
+        return self.router == ""
+
+    def __hash__(self):
+        return(hash(str(self)))
+    
+    def add_next_hop(self, other:"Hop"):
+        """
+        creates parent/child relationship between this hop and another one
+        """
+        self._children.append(other)
+        other._parent = self
+    
+
+
+class Path:
+    """
+    Keeps track of all our paths
+    """
+    def __init__(self, first_hop:Hop):
+
+        self.first_hop = first_hop
+        self.hops = {str(first_hop): first_hop}
+    
+    def add_hop(self, curr_hop:Hop, next_hop:Hop) -> bool:
+        """
+        Adds a hop to our current hop. If this is a new hop,
+        which means we're on a new path, returns true. Otherwise
+        returns false to indicate we don't have to walk this path
+        since it's a duplicate of another.
+        """
+
+        # if we've seen this hop already we will check for loops.
+        if str(next_hop) in self.hops:
+            new_path = False
+
+            # parent = self.curr_hop._parent
+            # while parent:
+            #     if parent == next_hop:
+            #         raise ValueError(f"Loop detected for {self.curr_hop} -> {next_hop}")
+            #     parent = parent._parent
+
+            self.hops[str(next_hop)]
+
+        else:
+            new_path = True
+            self.hops[str(next_hop)] = next_hop
+            
+        curr_hop.add_next_hop(next_hop)
+        return new_path
+    
+
+
+    def walk_path(self) -> Dict[int, List[Hop]]:
+        """
+        'Flattens' path from first hop to last into a dict
+        of list of hops keyed on hop number.
+        """
+        hop_num = 1
+        result = { hop_num: {self.first_hop}}
+        self._walk_path(self.first_hop, hop_num, result)
+        return result
+        
+            
+    def _walk_path(self, curr_hop, hop_num, result):
+        hop_num +=1
+        for child in curr_hop._children:
+            if hop_num not in result:
+                result[hop_num] = set()
+            result[hop_num].add(child)
+            self._walk_path(child, hop_num, result)
+
+ 
+    def print_path(self):
+        """
+        Prints out the path, hop by hop
+        """
+        for hop_num, hops in self.walk_path().items():
+            hop_list = list(hops)
+            hop_list.sort()
+            print(f"***** hop {hop_num} *****")
+            for hop in hop_list:
+                print(f"\t{hop}")
+
+
+    def get_path(self) -> list[dict]:
+        """
+        Returns hops as a list of dicts - for very basic displaying in
+        a table. In the future want to beef up the cli script
+        so that it can do nested tables, which is really
+        what we need for this.
+        """
+        result = []
+        for hop_num, hops in self.walk_path().items():
+            hop_list = list(hops)
+            hop_list.sort()
+            hop_str = "\n".join([str(h) for h in hop_list]) + "\n"              
+            result.append({'hop_num':hop_num, 'hops': hop_str})
+
+        return result

{umnetdb_utils-0.1.5.dist-info → umnetdb_utils-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: umnetdb-utils
-Version: 0.1.5
+Version: 0.2.0
 Summary: Helper classes for querying UMnet databases
 License: MIT
 Author: Amy Liebowitz
@@ -16,6 +16,7 @@ Requires-Dist: oracledb (>=3.1.0,<4.0.0)
 Requires-Dist: psycopg[binary] (>=3.2.9,<4.0.0)
 Requires-Dist: python-decouple (>=3.8,<4.0)
 Requires-Dist: sqlalchemy (>=2.0.41,<3.0.0)
+Requires-Dist: typer (>=0.16.0,<0.17.0)
 Description-Content-Type: text/markdown
 
 # umnetdb-utils

umnetdb_utils-0.2.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+umnetdb_utils/__init__.py,sha256=YrT7drJ2MlU287z4--8EQI3fKLQZBlVxaCacU1yoDtg,132
+umnetdb_utils/base.py,sha256=LWVhrbShhuYekXZxwvhhVUkH6DHd_wnJ2QZ_zw3fKqM,6139
+umnetdb_utils/cli.py,sha256=tq_XxGrCWJacP3tzxTULPHeuvbh1AU_GHXHY1jkI4-U,3349
+umnetdb_utils/umnetdb.py,sha256=bzMwItOfCgpKhCGwmghqAyp9gKvckGsUIGvSJC1iUho,17122
+umnetdb_utils/umnetdisco.py,sha256=D9MiF_QrYznYL_ozLUmD1ASzkGDs1jqZ_XksNV7wW3o,8018
+umnetdb_utils/umnetequip.py,sha256=Nnh_YCcqGpYRp2KsS8qbAGJ9AgTeYB4umeYeP9_B-6A,5440
+umnetdb_utils/umnetinfo.py,sha256=vqFvLbIu7Wu650RTqzaELpbKhq6y91fgBXM5Q_elRl4,18443
+umnetdb_utils/utils.py,sha256=kAVLjX4UvN5VdllYPN2t-q386iDjsZZ-Bp8G_ZTchbI,9418
+umnetdb_utils-0.2.0.dist-info/METADATA,sha256=kkv7rHDme5K2VXBfe24KFg6qBOsezgLZc8XYku5dpwU,1595
+umnetdb_utils-0.2.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+umnetdb_utils-0.2.0.dist-info/entry_points.txt,sha256=nNmxVRlYLaeepjDKNCyARBF3d8igHtZGTVA_zoglevI,50
+umnetdb_utils-0.2.0.dist-info/RECORD,,

umnetdb_utils-0.2.0.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[console_scripts]
+umnetdb=umnetdb_utils.cli:main
+

umnetdb_utils-0.1.5.dist-info/RECORD

@@ -1,10 +0,0 @@
-umnetdb_utils/__init__.py,sha256=YrT7drJ2MlU287z4--8EQI3fKLQZBlVxaCacU1yoDtg,132
-umnetdb_utils/base.py,sha256=LWVhrbShhuYekXZxwvhhVUkH6DHd_wnJ2QZ_zw3fKqM,6139
-umnetdb_utils/umnetdb.py,sha256=jHifl5eJXzye28mjLnfn7A7zwN2Y4YuKgrFvjV5MQBI,4953
-umnetdb_utils/umnetdisco.py,sha256=D9MiF_QrYznYL_ozLUmD1ASzkGDs1jqZ_XksNV7wW3o,8018
-umnetdb_utils/umnetequip.py,sha256=Nnh_YCcqGpYRp2KsS8qbAGJ9AgTeYB4umeYeP9_B-6A,5440
-umnetdb_utils/umnetinfo.py,sha256=vqFvLbIu7Wu650RTqzaELpbKhq6y91fgBXM5Q_elRl4,18443
-umnetdb_utils/utils.py,sha256=VLxfBKq1iOQVxTuySvxV4goE_jtX4jCshjf9a5ky1vI,989
-umnetdb_utils-0.1.5.dist-info/METADATA,sha256=fPpx67_VnM0Xp6JGk13mRrxOte8wPAkau0O0PtM9vW0,1555
-umnetdb_utils-0.1.5.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-umnetdb_utils-0.1.5.dist-info/RECORD,,