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