More spring cleaning.
[python_utils.git] / arper.py
1 #!/usr/bin/env python3
2
3 """A caching layer around the kernel's network mapping between IPs and MACs"""
4
5
6 import datetime
7 import logging
8 import os
9 import warnings
10 from typing import Any, Optional
11
12 from overrides import overrides
13
14 import argparse_utils
15 import config
16 import exec_utils
17 import file_utils
18 import persistent
19 import site_config
20 import string_utils
21 from collect.bidict import BiDict
22
23 logger = logging.getLogger(__name__)
24
25 cfg = config.add_commandline_args(
26     f'MAC <--> IP Address mapping table cache ({__file__})',
27     'Commandline args related to MAC <--> IP Address mapping',
28 )
29 cfg.add_argument(
30     '--arper_cache_location',
31     default=site_config.get_config().arper_cache_file,
32     metavar='FILENAME',
33     help='Where to cache the kernel ARP table',
34 )
35 cfg.add_argument(
36     '--arper_supplimental_cache_location',
37     default=site_config.get_config(site_config.other_location()).arper_cache_file,
38     metavar='FILENAME',
39     help='Where someone else is caching the kernel ARP table',
40 )
41 cfg.add_argument(
42     '--arper_cache_max_staleness',
43     type=argparse_utils.valid_duration,
44     default=datetime.timedelta(seconds=60 * 15),
45     metavar='DURATION',
46     help='Max acceptable age of the kernel arp table cache',
47 )
48 cfg.add_argument(
49     '--arper_min_entries_to_be_valid',
50     type=int,
51     default=site_config.get_config().arper_minimum_device_count,
52     help='Min number of arp entries to bother persisting.',
53 )
54
55
56 @persistent.persistent_autoloaded_singleton()  # type: ignore
57 class Arper(persistent.Persistent):
58     """A caching layer around the kernel's network mapping between IPs and
59     MACs.  This class restores persisted state that expires
60     periodically (see --arper_cache_max_staleness) at program startup
61     time.  If it's unable to use the file's contents, it queries the
62     kernel (via arp) and uses an auxillary utility called arp-scan to
63     query the network.  If it has to do this there's a latency hit but
64     it persists the collected data in the cache file.  Either way, the
65     class behaves as a global singleton hosting this data thereafter.
66
67     """
68
69     def __init__(
70         self,
71         cached_local_state: Optional[BiDict] = None,
72         cached_supplimental_state: Optional[BiDict] = None,
73     ) -> None:
74         self.state = BiDict()
75         if cached_local_state is not None:
76             logger.debug('Loading Arper map from cached local state.')
77             self.state = cached_local_state
78         else:
79             logger.debug('No usable cached state; calling /usr/sbin/arp')
80             self.update_from_arp_scan()
81             self.update_from_arp()
82         if len(self.state) < config.config['arper_min_entries_to_be_valid']:
83             raise Exception(f'Arper didn\'t find enough entries; only got {len(self.state)}.')
84         if cached_supplimental_state is not None:
85             logger.debug('Also added %d supplimental entries.', len(cached_supplimental_state))
86             for mac, ip in cached_supplimental_state.items():
87                 self.state[mac] = ip
88         for mac, ip in self.state.items():
89             logger.debug('%s <-> %s', mac, ip)
90
91     def update_from_arp_scan(self):
92         network_spec = site_config.get_config().network
93         try:
94             output = exec_utils.cmd(
95                 f'/usr/local/bin/arp-scan --retry=6 --timeout 350 --backoff=1.4 --random --numeric --plain --ignoredups {network_spec}',
96                 timeout_seconds=10.0,
97             )
98         except Exception as e:
99             logger.exception(e)
100             return
101         for line in output.split('\n'):
102             ip = string_utils.extract_ip_v4(line)
103             mac = string_utils.extract_mac_address(line)
104             if ip is not None and mac is not None and mac != 'UNKNOWN' and ip != 'UNKNOWN':
105                 mac = mac.lower()
106                 logger.debug('ARPER: %s => %s', mac, ip)
107                 self.state[mac] = ip
108
109     def update_from_arp(self):
110         try:
111             output = exec_utils.cmd('/usr/sbin/arp -a', timeout_seconds=10.0)
112         except Exception as e:
113             logger.exception(e)
114             return
115         for line in output.split('\n'):
116             ip = string_utils.extract_ip_v4(line)
117             mac = string_utils.extract_mac_address(line)
118             if ip is not None and mac is not None and mac != 'UNKNOWN' and ip != 'UNKNOWN':
119                 mac = mac.lower()
120                 logger.debug('ARPER: %s => %s', mac, ip)
121                 self.state[mac] = ip
122
123     def get_ip_by_mac(self, mac: str) -> Optional[str]:
124         mac = mac.lower()
125         return self.state.get(mac, None)
126
127     def get_mac_by_ip(self, ip: str) -> Optional[str]:
128         return self.state.inverse.get(ip, None)
129
130     @classmethod
131     def load_state(
132         cls,
133         cache_file: str,
134         freshness_threshold_sec: int,
135         state: BiDict,
136     ):
137         if not file_utils.file_is_readable(cache_file):
138             logger.debug('Can\'t read %s', cache_file)
139             return
140         if persistent.was_file_written_within_n_seconds(
141             cache_file,
142             freshness_threshold_sec,
143         ):
144             logger.debug('Loading state from %s', cache_file)
145             count = 0
146             with open(cache_file, 'r') as rf:
147                 contents = rf.readlines()
148                 for line in contents:
149                     line = line[:-1]
150                     logger.debug('ARPER:%s> %s', cache_file, line)
151                     (mac, ip) = line.split(',')
152                     mac = mac.strip()
153                     mac = mac.lower()
154                     ip = ip.strip()
155                     state[mac] = ip
156                     count += 1
157         else:
158             logger.debug('%s is too stale.', cache_file)
159
160     @classmethod
161     @overrides
162     def load(cls) -> Any:
163         local_state: BiDict = BiDict()
164         cache_file = config.config['arper_cache_location']
165         max_staleness = config.config['arper_cache_max_staleness'].total_seconds()
166         logger.debug('Trying to load main arper cache from %s...', cache_file)
167         cls.load_state(cache_file, max_staleness, local_state)
168         if len(local_state) <= config.config['arper_min_entries_to_be_valid']:
169             msg = f'{cache_file} is invalid: only {len(local_state)} entries.  Deleting it.'
170             logger.warning(msg)
171             warnings.warn(msg, stacklevel=2)
172             try:
173                 os.remove(cache_file)
174             except Exception:
175                 pass
176
177         supplimental_state: BiDict = BiDict()
178         cache_file = config.config['arper_supplimental_cache_location']
179         max_staleness = config.config['arper_cache_max_staleness'].total_seconds()
180         logger.debug('Trying to suppliment arper state from %s', cache_file)
181         cls.load_state(cache_file, max_staleness, supplimental_state)
182         if len(local_state) > 0:
183             return cls(local_state, supplimental_state)
184         return None
185
186     @overrides
187     def save(self) -> bool:
188         if len(self.state) > config.config['arper_min_entries_to_be_valid']:
189             logger.debug('Persisting state to %s', config.config["arper_cache_location"])
190             with file_utils.FileWriter(config.config['arper_cache_location']) as wf:
191                 for (mac, ip) in self.state.items():
192                     mac = mac.lower()
193                     print(f'{mac}, {ip}', file=wf)
194             return True
195         else:
196             logger.warning(
197                 'Only saw %d entries; needed at least %d to bother persisting.',
198                 len(self.state),
199                 config.config["arper_min_entries_to_be_valid"],
200             )
201             return False