nbdig

#!/usr/bin/env python3
# nbdig (part of ossobv/vcutil) // wdoekes/2024 // Public Domain
#
# Quickly query the netbox server for hosts and hostnames/IPs, using the
# Netbox REST API.
#
# Usage examples:
#
#     $ nbdig walter  # or 'w?lter.*'
#     walter.example.com              123.123.123.123
#     walter.internal.lan             10.32.1.5
#
#     $ nbdig walter.internal.lan
#     10.32.1.5
#
#     $ nbdig -x 10.32.1.5
#     walter.internal.lan
#
#     $ nbdig -x 10.32.1.0/24 -t
#     10.32.0.0/22                    ROOT
#       10.32.1.0/24                  container /24
#         10.32.1.0/27                usable /31 space
#           10.32.1.4/31              walter.internal.lan
#
# Configure by setting up a ~/.config/nbdig.ini with one section:
#
#     [netbox.example.com]
#     api_url = https://netbox.example.com/api
#     api_token = 0123456789abcdef0123456789abcdef01234567
#
from argparse import ArgumentParser
from collections import namedtuple
from configparser import ConfigParser, MissingSectionHeaderError
from fnmatch import translate as glob_to_re
from ipaddress import IPv4Address, IPv4Network
from json import dumps
from os import path
from re import IGNORECASE, compile as re_compile, split as re_split
from sys import stderr
from warnings import warn

import requests
import warnings


INI_DEFAULT = '~/.config/nbdig.ini'


def net2ip(s):
    "'1.2.3.4/31' -> IPv4Address('1.2.3.4')"
    return IPv4Address(s.rsplit('/', 1)[0])


class StartupError(ValueError):
    pass


class Config(namedtuple('Config', 'api_url api_token')):
    @classmethod
    def from_defaults(cls):
        return cls.from_ini(path.expanduser(INI_DEFAULT))

    @classmethod
    def from_ini(cls, filename):
        ini = ConfigParser(
            delimiters=('=',), allow_no_value=True,
            inline_comment_prefixes=('#', ';'))
        try:
            ini.read(filename)
        except MissingSectionHeaderError as e:
            raise StartupError(f'{e} in in {filename}') from e
        assert len(ini.sections()) == 1, ini.sections()
        the_section = ini.sections()[0]
        data = dict(ini.items(the_section))
        try:
            # api_url, with "/api": "https://netbox.example.com/api"
            return cls(api_url=data['api_url'], api_token=data['api_token'])
        except KeyError as e:
            raise StartupError(
                f'api_url or api_token not found in {filename}') from e


class IpAddress(namedtuple('IpAddress', 'address name iface vrf is_primary')):
    def __new__(cls, *args, **kwargs):
        if 'is_primary' not in kwargs:
            kwargs['is_primary'] = False
        return super().__new__(cls, *args, **kwargs)

    def as_primary(self):
        return IpAddress(
            address=self.address, name=self.name,
            iface=self.iface, vrf=self.vrf,
            is_primary=True)

    @property
    def optname(self):
        if not self.name:
            return ''
        return '"{}"'.format(self.name.replace('"', '""'))

    @property
    def safeiface(self):
        if not self.iface:
            return '-'
        if self.iface.startswith('GigabitEthernet'):
            return f'GE{self.iface[15:]}'
        return self.iface

    @property
    def safevrf(self):
        return self.vrf or '-'


class Host(namedtuple('Host', 'name ip_addresses details')):
    def __new__(cls, *args, **kwargs):
        if 'details' not in kwargs:
            kwargs['details'] = None
        return super().__new__(cls, *args, **kwargs)

    @property
    def optdetails(self):
        if not self.details:
            return ''
        return '{{{}}}'.format(
            self.details.replace('{', '').replace('}', ''))

    @property
    def safename(self):
        return self.name or '<null>'  # or '-'?


class Network(namedtuple('Network', 'prefix description children')):
    def merge(self, subnets):
        if not self.children:
            assert all(s.prefix.subnet_of(self.prefix) for s in subnets), (
                self, subnets)
            self.children.extend(subnets)
            return

        assert len(self.children) == 1, self.children
        self.children[0].merge(subnets)

    def sort(self):
        for child in self.children:
            child.sort()
        self.children.sort()


class ThinNetboxWrapper:
    NB_DEVS = '/dcim/devices/'
    NB_IPADDR = '/ipam/ip-addresses/'
    NB_PREFIX = '/ipam/prefixes/'
    NB_VMS = '/virtualization/virtual-machines/'

    def __init__(self, config):
        self._config = config

    def _get_query(self, path, filters):
        url = f'{self._config.api_url}{path}'
        headers = {'Authorization': f'Token {self._config.api_token}'}
        resp = requests.get(url, params=filters, headers=headers)
        resp.raise_for_status()
        return resp

    def find_hostname(self, hostfragment, load_interfaces, load_properties):
        if '*' in hostfragment or '?' in hostfragment:
            params = [
                ('name__ic', part)
                for part in re_split('[*?]+', hostfragment) if part]
            # Alas. Multilpe name-searches will do an OR search. Take the
            # largest string and hope it's the most narrow match.
            params.sort(key=(lambda x: (-len(x[1]), x[1])))
            params = params[0:1]
            # Turn into a regex.
            re_match = re_compile(glob_to_re(hostfragment), IGNORECASE)
        else:
            params = [('name__isw', hostfragment)]  # startswith!
            # Turn into a regex. Make it a startswith match for now.
            re_match = re_compile(glob_to_re(hostfragment + '*'), IGNORECASE)

        assert load_interfaces in (False, True), load_interfaces

        show_properties = []
        load_properties = load_properties or ()

        for prop in load_properties:
            # See API docs:
            # https://demo.netbox.dev/static/docs/rest-api/filtering/
            if '=' in prop:
                prop_lookup, prop_value = prop.split('=', 1)
                if prop_value and prop_value[0] not in (
                        'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
                        'abcdefghijklmnopqrstuvwxyz'
                        '0123456789_'):
                    if prop_value[0] == '^' and len(prop_value) > 1:
                        params.append((f'{prop_lookup}__isw', prop_value[1:]))
                    elif prop_value[0] == '^':
                        # Exclude "exactly empty", so include non-empty.
                        params.append((f'{prop_lookup}__nie', ''))
                        # For numeric fields, including only 0 and above.
                        params.append((f'{prop_lookup}__gte', '0'))
                    else:
                        raise NotImplementedError(
                            'known =-modifiers are: =EQUALS =^STARTSWITH '
                            '(use =^ to match for non-empty)')
                else:
                    # Do not use '{prop_lookup}__ie' here: it will break
                    # fancy value lookup, like status=offline, where
                    # status is a dict.
                    params.append((prop_lookup, prop_value))
            else:
                prop_lookup = prop

            if prop_lookup not in show_properties:
                show_properties.append(prop_lookup)

        # We'll need to get all the info in a second run. For large resultsets
        # we want to trim them by globbing name first.
        params.extend([('brief', 1), ('limit', 1000)])

        # Get device/VM ids.
        dev_ids, dev_ids_x = self._search_for_ids(
            self.NB_DEVS, params, hostfragment, re_match)
        vm_ids, vm_ids_x = self._search_for_ids(
            self.NB_VMS, params, hostfragment, re_match)

        # Are there any exact matches?
        if dev_ids_x or vm_ids_x:
            # Then take only exact matches.
            dev_ids = dev_ids_x
            vm_ids = vm_ids_x

        # Collect full info on the selected hosts.
        hosts = []
        for pth, lookup_key, ids in (
                (self.NB_DEVS, 'device_id', dev_ids),
                (self.NB_VMS, 'virtual_machine_id', vm_ids)):
            if ids:
                # Get interface infos.
                if load_interfaces:
                    interfaces_by_id = self._find_interfaces(lookup_key, ids)
                else:
                    interfaces_by_id = {}

                # Get host infos.
                params = [('id', id_) for id_ in ids]
                params.extend([('limit', 1000), ('exclude', 'config_context')])
                ret = self._get_query(pth, params)
                pages = ret.json()
                assert pages['next'] is None, NotImplemented  # pagination
                for dev_or_vm in pages['results']:
                    hosts.append(self._nb_device_or_vm_to_host(
                        dev_or_vm, interfaces_by_id=interfaces_by_id,
                        show_properties=show_properties))

        # Sort, and cope with Nones.
        hosts.sort(key=(lambda x: (x.name or '')))

        return hosts

    def _nb_device_or_vm_to_host(
            self, info, interfaces_by_id, show_properties):

        name = info['name']
        ip_addresses = []

        primary_ip = info['primary_ip']
        if primary_ip is not None:
            primary_ip = IpAddress(
                address=net2ip(primary_ip['address']), name='?',  # with '?'?
                iface='?', vrf='?', is_primary=True)

        if info['id'] in interfaces_by_id:
            ip_addresses.extend(interfaces_by_id[info['id']])
            if primary_ip is not None:
                for idx, ip in enumerate(ip_addresses):
                    if ip.address == primary_ip.address:
                        ip_addresses[idx] = ip.as_primary()
                        break
        elif primary_ip is not None:
            ip_addresses.append(primary_ip)
        else:
            warn(f'{name}: no primary IP set')

        def kv(k, v):
            if isinstance(v, dict):
                if 'slug' in v:
                    v = v['slug']
                elif 'name' in v:
                    v = v['name']
                elif 'id' in v:
                    v = v['id']
                else:
                    assert False, v
            return '{}={}'.format(k, v or '')

        details = ' '.join(kv(k, v) for k, v in self._nb_extract_properties(
            info, show_properties))

        return Host(name=name, ip_addresses=ip_addresses, details=details)

    def _find_interfaces(self, lookup_key, ids):
        assert lookup_key in ('device_id', 'virtual_machine_id'), lookup_key
        lookup_short = lookup_key[0:-3]  # 'device', 'virtual_machine'

        params = [(lookup_key, id_) for id_ in ids]
        params.append(('limit', 1000))
        ret = self._get_query(self.NB_IPADDR, params)
        pages = ret.json()
        assert pages['next'] is None, NotImplemented  # pagination

        interfaces_by_id = {}
        for interface in pages['results']:
            assert interface['family'] == {'value': 4, 'label': 'IPv4'}, (
                interface['id'], interface['address'], interface['family'])
            assigned_address = net2ip(interface['address'])
            assigned_vrf = (interface.get('vrf') or {}).get('rd', None)
            assigned_to_id = interface['assigned_object'][lookup_short]['id']
            assigned_iface = interface['assigned_object']['name']

            if assigned_to_id not in interfaces_by_id:
                interfaces_by_id[assigned_to_id] = []
            assigned_name = (interface['dns_name'] or interface['description'])
            interfaces_by_id[assigned_to_id].append(
                IpAddress(
                    address=assigned_address, name=assigned_name,
                    iface=assigned_iface, vrf=assigned_vrf))

        return interfaces_by_id

    def _nb_extract_properties(self, info, show_properties):
        props = []
        for prop in show_properties:
            if prop.startswith('cf_'):
                prop_lookup = prop[3:]
                dict_ = info['custom_fields']
            else:
                prop_lookup = prop
                dict_ = info

            if prop_lookup in dict_:
                # Some properties are a dictionary. Take the "machine" value.
                if isinstance(dict_[prop_lookup], dict) and (
                        'value' in dict_[prop_lookup]):
                    props.append((prop, dict_[prop_lookup]['value']))
                elif isinstance(dict_[prop_lookup], list) and any(
                        'slug' in v for v in dict_[prop_lookup]):
                    props.append((prop, '+'.join(
                        v['slug'] for v in dict_[prop_lookup])))
                else:
                    props.append((prop, dict_[prop_lookup]))
            else:
                # FIXME: This warning might go off for properties that
                # exist on devices but not on VMs; e.g. for 'rear_port_count'.
                others = [k for k in info.keys() if k != 'custom_fields'] + [
                    f'cf_{k}' for k in info['custom_fields']]
                others = ', '.join(others)
                warn(f'property {prop!r} not found. maybe: {others}')

        return props

    def _search_for_ids(self, path, filters, exact_match, regex_match):
        exact_match = exact_match.lower()

        ret = self._get_query(path, filters)
        pages = ret.json()
        assert pages['next'] is None, NotImplemented  # pagination
        ids_exact = [
            res['id'] for res in pages['results']
            if (res.get('name') or '').lower() == exact_match]
        ids = [
            res['id'] for res in pages['results']
            if regex_match.match(res.get('name') or '')]

        return ids, ids_exact

    def find_ip_address(self, ip_address):
        cidr = IPv4Network(ip_address)
        if cidr.prefixlen != 32:
            params = [('limit', 1000), ('parent', str(cidr))]
        else:
            params = [('limit', 1000), ('address', str(cidr.network_address))]
        ret = self._get_query(self.NB_IPADDR, params)
        pages = ret.json()
        assert pages['next'] is None, NotImplemented  # pagination
        try:
            hosts = self._find_ip_address(pages)
        except KeyError as e:
            print(f'error: {e}', file=stderr)
            print('input: {}'.format(dumps(pages)), file=stderr)
            raise

        return hosts

    def _find_ip_address(self, pages):
        results = pages['results']
        assigned = []
        for address in results:
            # Discard network bits for now
            ip_address = net2ip(address['address'])
            vrf = (address.get('vrf') or {}).get('rd', None)
            assignment = address['assigned_object']
            iface = (assignment['name'] if assignment else None)

            if assignment:
                if 'virtual_machine' in assignment:
                    name = assignment['virtual_machine']['name']
                elif 'device' in assignment:
                    name = assignment['device']['name']
                else:
                    raise NotImplementedError(assignment)
                ip_name = address['dns_name'] or address['description']
            elif address['dns_name'] and address['description']:
                name, ip_name = address['dns_name'], address['description']
            else:
                name = address['dns_name'] or None
                ip_name = address['dns_name'] or address['description'] or ''

            ipaddr = IpAddress(
                address=ip_address, name=ip_name, iface=iface, vrf=vrf)
            host = Host(name=name, ip_addresses=[ipaddr])
            assigned.append(host)

        if pages['count'] != len(assigned):
            raise ValueError('unhandled count: {}, got {}'.format(
                pages, assigned))

        return assigned

    def find_networks(self, cidr):
        cidr = IPv4Network(cidr)
        params = [('limit', 1000), ('contains', str(cidr))]
        ret = self._get_query(self.NB_PREFIX, params)
        pages = ret.json()
        assert pages['next'] is None, NotImplemented  # pagination
        try:
            supernets = self._find_networks(pages)
        except KeyError as e:
            print(f'error: {e}', file=stderr)
            print('input: {}'.format(dumps(pages)), file=stderr)
            raise

        params = [('limit', 1000), ('within', str(cidr))]
        ret = self._get_query(self.NB_PREFIX, params)
        pages = ret.json()
        assert pages['next'] is None, NotImplemented  # pagination
        try:
            subnets = self._find_networks(pages)
        except KeyError as e:
            print(f'error: {e}', file=stderr)
            print('input: {}'.format(dumps(pages)), file=stderr)
            raise

        # Merge supernets and subnets.
        if supernets and subnets:
            assert len(supernets) == 1 and subnets, (
                supernets, subnets)
            supernets[0].merge(subnets)
        elif supernets:
            pass
        elif subnets:
            supernets = subnets

        return supernets

    def _find_networks(self, pages):
        results = pages['results']
        nets = []
        for res in results:
            net = Network(
                prefix=IPv4Network(res['prefix']),
                description=res['description'],
                children=[])
            nets.append(net)

        # Sort networks: larger first.
        nets.sort(key=(lambda x: x.prefix.prefixlen))

        # Move smaller networks into larger ones if possible.
        i = len(nets) - 1
        while i > 0:
            small_net = nets[i]
            for j in range(i - 1, -1, -1):
                larger_net = nets[j]
                if larger_net.prefix.prefixlen > small_net.prefix.prefixlen:
                    break  # quit early
                if small_net.prefix.subnet_of(larger_net.prefix):
                    larger_net.children.append(small_net)
                    nets.pop(i)
                    break
            i -= 1

        # nets is now a recursive structure. Sort it, this time by prefix.
        [net.sort() for net in nets]    # let the Network()s sort themselves
        nets.sort()                     # sort the Network()s

        return nets


def show_ipaddresses(hosts, interface_match=None):
    # Concatenate to a big list.
    ip_addresses = []
    for host in hosts:
        ip_addresses.extend(host.ip_addresses)

    # Sort. We want primary first.
    ip_addresses = sorted(ip_addresses, key=(lambda x: (
        not x.is_primary, x.address, x.vrf, x.iface)))

    ip_addresses = [
        str(ip.address) for ip in ip_addresses
        if interface_match is None or interface_match in ip.iface]

    print('\n'.join(ip_addresses))


def show_name_addr_extra(hosts):
    for host in hosts:
        addr = ' '.join(str(ip.address) for ip in host.ip_addresses) or '-'
        print(f'{host.safename!s:30}  {addr!s:15}  {host.optdetails}'.rstrip())


def show_name_addr_iface_vrf(hosts, interface_match=None):
    for host in hosts:
        host_name = host.safename
        host_details = host.optdetails

        # Sort. We want primary first.
        ip_addresses = sorted(host.ip_addresses, key=(lambda x: (
            not x.is_primary, x.address, x.vrf, x.iface)))

        for ip in ip_addresses:
            if interface_match is None or interface_match in ip.iface:
                addr = (
                    f'[{ip.address}]' if ip.is_primary
                    else f' {ip.address} ')

                if host_details:
                    extra = host_details
                else:
                    extra = ip.optname

                print(
                    f'{host_name!s:30} {addr!s:17} {ip.safeiface!s:15}  '
                    f'{ip.safevrf!s:15}  {extra}'.rstrip())


def do_interface_lookup(
        nb, hostfragment, interface_match, load_properties, verbose=False):

    if load_properties:
        raise NotImplementedError('not sure what to do with properties here')
    load_properties = ()

    hosts = nb.find_hostname(
        hostfragment, load_interfaces=True, load_properties=load_properties)
    if not hosts:
        print('not found', file=stderr)
        return

    if verbose:
        show_name_addr_iface_vrf(hosts, interface_match)
    else:
        show_ipaddresses(hosts, interface_match)


def do_forward_lookup(nb, hostfragment, load_properties, verbose=False):
    hosts = nb.find_hostname(
        hostfragment, load_interfaces=False, load_properties=load_properties)

    if not hosts:
        print('not found', file=stderr)
        return

    if not (len(hosts) == 1 and (
            # exact match?
            hostfragment == hosts[0].name or
            # intentional * for exact match?
            '*' in hostfragment or '?' in hostfragment)):
        verbose = True

    if verbose:
        show_name_addr_extra(hosts)
    else:
        # There is only one host, but the function functions the same.
        show_ipaddresses(hosts)


def do_reverse_lookup(nb, ip_or_cidr, tree=False, verbose=False):
    if not tree:
        hosts = nb.find_ip_address(ip_or_cidr)
        networks = None
    elif tree and not verbose:
        hosts = None
        networks = nb.find_networks(ip_or_cidr)
    else:
        raise NotImplementedError(
            'lookup with networks and multiple hosts not implemented')

    if not hosts and not networks:
        print('not found', file=stderr)
        return

    if verbose:
        assert hosts and not networks
        show_name_addr_iface_vrf(hosts)
    elif hosts:
        assert not networks
        all_hostnames = set(h.safename for h in hosts)
        print('\n'.join(sorted(all_hostnames)))
    else:
        def print_nets(nets, indent=''):
            for net in nets:
                indent_addr = f'{indent}{net.prefix}'
                print(f'{indent_addr!s:30}  {net.description}')
                print_nets(net.children, indent=(indent + '  '))

        print_nets(networks)


def formatwarning(message, category, filename, lineno, line=None):
    """
    Override default Warning layout, from:

        /usr/bin/nbdig:174: UserWarning: DD6-B FREE: no primary IP set
            warn(f'{name}: no primary IP set')

    To:

        nbdig:175: UserWarning: DD6-B FREE: no primary IP set
    """
    a = b = ''
    if stderr.isatty():
        a, b = '\x1b[31;1m', '\x1b[0m'
    return '{a}{basename}:{lineno}: {category}: {message}{b}\n'.format(
        a=a, b=b, basename=path.basename(filename), lineno=lineno,
        category=category.__name__, message=message)


def main():
    warnings.formatwarning = formatwarning  # noqa

    parser = ArgumentParser(prog='nbdig', description='(zab)dig for netbox')
    parser.add_argument(
        '-c', '--config', metavar='INIFILE',
        help=f'configuration INI location (default: {INI_DEFAULT})')
    parser.add_argument(
        '-i', '--iface', metavar='IFACE', help='interface lookup')
    parser.add_argument(
        '-p', '--property', action='append', metavar='PROPERTY',
        help=('extra properties to search/list for forward lookup; '
              'e.g. status, or status=offline to filter directly; or '
              'a custom field cf_machine_id=^2 that starts with a 2'))
    parser.add_argument(
        '-t', '--tree', action='store_true', help='show network tree for -x')
    parser.add_argument('-x', metavar='NET', help='reverse lookup (IP or net)')
    parser.add_argument('-v', '--verbose', action='store_true', help='verbose')
    parser.add_argument('hostfragment', nargs='?')

    args = parser.parse_args()
    if args.x is not None and args.hostfragment is not None:
        parser.error("supply either 'hostfragment' or -x 'IP', not both")
    elif args.x is args.hostfragment is None:
        parser.error("requires at least one argument")

    try:
        if args.config is None:
            config = Config.from_defaults()
        else:
            config = Config.from_ini(args.config)
    except StartupError as e:
        parser.error(str(e))

    nb = ThinNetboxWrapper(config)
    if args.hostfragment is not None:
        if args.iface is None:
            do_forward_lookup(
                nb, args.hostfragment, args.property, verbose=args.verbose)
        else:
            do_interface_lookup(
                nb, args.hostfragment, args.iface, args.property,
                verbose=args.verbose)
    elif args.x is not None:
        if args.property:
            raise NotImplementedError('-x with -p not implemented')
        do_reverse_lookup(nb, args.x, tree=args.tree, verbose=args.verbose)
    else:
        raise NotImplementedError('cannot get here')


if __name__ == '__main__':
    main()