arper.py

   1 #!/usr/bin/env python3
   2
   3 """A caching layer around the kernel's network mapping between IPs and MACs"""
   4
   5 import datetime
   6 import logging
   7 import os
   8 from typing import Any, Optional
   9 import warnings
  10
  11 from overrides import overrides
  12
  13 import argparse_utils
  14 from collect.bidict import BiDict
  15 import config
  16 import exec_utils
  17 import file_utils
  18 import persistent
  19 import string_utils
  20 import site_config
  21
  22 logger = logging.getLogger(__name__)
  23
  24 cfg = config.add_commandline_args(
  25     f'MAC <--> IP Address mapping table cache ({__file__})',
  26     'Commandline args related to MAC <--> IP Address mapping',
  27 )
  28 cfg.add_argument(
  29     '--arper_cache_location',
  30     default=f'{os.environ["HOME"]}/cache/.arp_table_cache',
  31     metavar='FILENAME',
  32     help='Where to cache the kernel ARP table',
  33 )
  34 cfg.add_argument(
  35     '--arper_cache_max_staleness',
  36     type=argparse_utils.valid_duration,
  37     default=datetime.timedelta(seconds=60 * 15),
  38     metavar='DURATION',
  39     help='Max acceptable age of the kernel arp table cache'
  40 )
  41 cfg.add_argument(
  42     '--arper_min_entries_to_be_valid',
  43     type=int,
  44     default=site_config.get_config().arper_minimum_device_count,
  45     help='Min number of arp entries to bother persisting.'
  46 )
  47
  48
  49 @persistent.persistent_autoloaded_singleton()
  50 class Arper(persistent.Persistent):
  51     def __init__(
  52             self, cached_state: Optional[BiDict] = None
  53     ) -> None:
  54         self.state = BiDict()
  55         if cached_state is not None:
  56             logger.debug('Loading Arper map from cached state.')
  57             self.state = cached_state
  58         else:
  59             logger.debug('No usable cached state; calling /usr/sbin/arp')
  60             self.update_from_arp_scan()
  61             self.update_from_arp()
  62         if len(self.state) < config.config['arper_min_entries_to_be_valid']:
  63             raise Exception(f'Arper didn\'t find enough entries; only got {len(self.state)}.')
  64
  65     def update_from_arp_scan(self):
  66         network_spec = site_config.get_config().network
  67         try:
  68             output = exec_utils.cmd(
  69                 f'/usr/local/bin/arp-scan --retry=6 --timeout 350 --backoff=1.4 --random --numeric --plain --ignoredups {network_spec}',
  70                 timeout_seconds=10.0
  71             )
  72         except Exception as e:
  73             logger.exception(e)
  74             return
  75         for line in output.split('\n'):
  76             ip = string_utils.extract_ip_v4(line)
  77             mac = string_utils.extract_mac_address(line)
  78             if ip is not None and mac is not None and mac != 'UNKNOWN' and ip != 'UNKNOWN':
  79                 mac = mac.lower()
  80                 logger.debug(f'ARPER: {mac} => {ip}')
  81                 self.state[mac] = ip
  82
  83     def update_from_arp(self):
  84         try:
  85             output = exec_utils.cmd(
  86                 '/usr/sbin/arp -a',
  87                 timeout_seconds=10.0
  88             )
  89         except Exception as e:
  90             logger.exception(e)
  91             return
  92         for line in output.split('\n'):
  93             ip = string_utils.extract_ip_v4(line)
  94             mac = string_utils.extract_mac_address(line)
  95             if ip is not None and mac is not None and mac != 'UNKNOWN' and ip != 'UNKNOWN':
  96                 mac = mac.lower()
  97                 logger.debug(f'ARPER: {mac} => {ip}')
  98                 self.state[mac] = ip
  99
 100     def get_ip_by_mac(self, mac: str) -> Optional[str]:
 101         mac = mac.lower()
 102         return self.state.get(mac, None)
 103
 104     def get_mac_by_ip(self, ip: str) -> Optional[str]:
 105         return self.state.inverse.get(ip, None)
 106
 107     @classmethod
 108     @overrides
 109     def load(cls) -> Any:
 110         cache_file = config.config['arper_cache_location']
 111         if persistent.was_file_written_within_n_seconds(
 112                 cache_file,
 113                 config.config['arper_cache_max_staleness'].total_seconds(),
 114         ):
 115             logger.debug(f'Loading state from {cache_file}')
 116             cached_state = BiDict()
 117             with open(cache_file, 'r') as rf:
 118                 contents = rf.readlines()
 119                 for line in contents:
 120                     line = line[:-1]
 121                     logger.debug(f'ARPER:{cache_file}> {line}')
 122                     (mac, ip) = line.split(',')
 123                     mac = mac.strip()
 124                     mac = mac.lower()
 125                     ip = ip.strip()
 126                     cached_state[mac] = ip
 127             if len(cached_state) > config.config['arper_min_entries_to_be_valid']:
 128                 return cls(cached_state)
 129             else:
 130                 msg = f'{cache_file} is invalid: only {len(cached_state)} entries.  Deleting it.'
 131                 logger.warning(msg)
 132                 warnings.warn(msg, stacklevel=2)
 133                 os.remove(cache_file)
 134         logger.debug('No usable saved state found')
 135         return None
 136
 137     @overrides
 138     def save(self) -> bool:
 139         if len(self.state) > config.config['arper_min_entries_to_be_valid']:
 140             logger.debug(
 141                 f'Persisting state to {config.config["arper_cache_location"]}'
 142             )
 143             with file_utils.FileWriter(config.config['arper_cache_location']) as wf:
 144                 for (mac, ip) in self.state.items():
 145                     mac = mac.lower()
 146                     print(f'{mac}, {ip}', file=wf)
 147             return True
 148         else:
 149             logger.warning(
 150                 f'Only saw {len(self.state)} entries; needed at least {config.config["arper_min_entries_to_be_valid"]} to bother persisting.'
 151             )
 152             return False