host: Add tools for discovering hosts

src/rouse/host.py

@@ -0,0 +1,209 @@
+from __future__ import unicode_literals
+import socket
+import struct
+import subprocess
+
+
+try:
+    long_ = long
+except NameError:
+    long_ = int
+try:
+    range_ = xrange
+except NameError:
+    range_ = range
+
+
+class Neighbor(object):
+
+    def __init__(self, ipaddr=None, macaddr=None):
+        self.ipaddr = ipaddr
+        self.macaddr = macaddr
+
+    def __repr__(self):
+        return str('{klass}({ipaddr!r}, {macaddr!r})'.format(
+            klass=self.__class__.__name__,
+            ipaddr=self.ipaddr,
+            macaddr=self.macaddr,
+        ))
+
+
+class NeighborTable(object):
+
+    def __init__(self):
+        self.entries = []
+        self.refresh()
+
+    def __iter__(self):
+        for n in self.entries:
+            yield n
+
+    def __getitem__(self, item):
+        for n in self:
+            if item in (n.ipaddr, n.macaddr):
+                return n
+        raise KeyError(item)
+
+    def refresh(self):
+        self.entries = []
+        try:
+            o = subprocess.check_output(['ip', 'neighbor', 'show'])
+        except (OSError, subprocess.CalledProcessError):
+            return
+        for line in reversed(o.splitlines()):
+            parts = line.strip().decode().split()
+            try:
+                i = parts[0]
+                m = MacAddress(parts[4]).with_colons()
+            except (IndexError, ValueError):
+                continue
+            self.entries.append(Neighbor(i, m))
+
+
+class MacAddress(object):
+
+    def __init__(self, macaddr):
+        if hasattr(macaddr, 'encode'):
+            self.value = self.parse(macaddr)
+        elif hasattr(macaddr, 'value'):
+            self.value = macaddr.value
+        else:
+            self.value = long_(macaddr)
+
+    def __str__(self):
+        return str(self.with_colons())
+
+    @classmethod
+    def parse(cls, macaddr):
+        if '.' in macaddr:
+            parts = cls.parse_dotted(macaddr)
+        elif ':' in macaddr:
+            parts = cls.parse_colons(macaddr)
+        elif '-' in macaddr:
+            parts = cls.parse_hyphens(macaddr)
+        else:
+            parts = (macaddr[i:i+2] for i in range_(0, len(macaddr), 2))
+        value = 0
+        for part in parts:
+            value <<= 8
+            value += int(part, 16)
+        return value
+
+    @classmethod
+    def parse_colons(cls, macaddr):
+        return macaddr.split(':')
+
+    @classmethod
+    def parse_dotted(cls, macaddr):
+        parts = macaddr.split('.')
+        if len(parts) != 3:
+            raise ValueError('Invalid MAC address {}'.format(macaddr))
+        for part in parts:
+            yield part[:2]
+            yield part[2:]
+
+    @classmethod
+    def parse_hyphens(cls, macaddr):
+        return macaddr.split('-')
+
+    @property
+    def _packed(self):
+        return bytearray(struct.pack(str('>Q'), self.value))[2:]
+
+    def with_colons(self):
+        return ':'.join('{:02x}'.format(x) for x in self._packed)
+
+    def with_dots(self):
+        p = list(self._packed)
+        parts = (p[i:i+2] for i in range_(0, len(p), 2))
+        return '.'.join('{:04x}'.format((x << 8) + y) for x, y in parts)
+
+    def with_hyphens(self):
+        return '-'.join('{:02X}'.format(x) for x in self._packed)
+
+
+class Host(object):
+
+    @staticmethod
+    def ping(host, port=7):
+        try:
+            ai = socket.getaddrinfo(host, port, socket.AF_UNSPEC,
+                                    socket.SOCK_DGRAM)
+        except socket.gaierror:
+            return
+        for af, socktype, proto, __, sa in ai:
+            sock = None
+            try:
+                sock = socket.socket(af, socktype, proto)
+                sock.sendto(b'\x00', sa)
+            except socket.error:
+                return
+            finally:
+                if sock is not None:
+                    sock.close()
+
+    @classmethod
+    def discover(cls, hint):
+        def try_mac():
+            try:
+                return cls.from_macaddr(hint)
+            except ValueError:
+                return None
+
+        def try_name():
+            try:
+                return cls.from_hostname(hint)
+            except ValueError:
+                return None
+
+        if ':' in hint:
+            return try_mac() or try_name()
+        else:
+            return try_name() or try_mac()
+
+    @classmethod
+    def from_hostname(cls, value):
+        try:
+            ai = socket.getaddrinfo(value, None, 0, 0, 0, socket.AI_CANONNAME)
+        except socket.gaierror as e:
+            raise ValueError(str(e))
+        ipaddr = None
+        hostname = None
+        for __, __, __, name, addr in ai:
+            if addr[0] != '::1' and not addr[0].startswith('127.'):
+                ipaddr = addr[0]
+            if name and not name.startswith('localhost'):
+                hostname = name
+            if hostname and ipaddr:
+                break
+        if not ipaddr:
+            return
+        cls.ping(ipaddr)
+        try:
+            neighbor = NeighborTable()[ipaddr]
+        except KeyError:
+            return None
+        host = cls()
+        host.macaddr = neighbor.macaddr
+        host.ipaddr = ipaddr
+        host.hostname = hostname
+        return host
+
+    @classmethod
+    def from_macaddr(cls, value):
+        m = MacAddress(value)
+        host = cls()
+        host.macaddr = m.with_colons()
+        try:
+            neighbor = NeighborTable()[host.macaddr]
+        except KeyError:
+            pass
+        else:
+            host.ipaddr = neighbor.ipaddr
+            try:
+                ni = socket.getnameinfo((host.ipaddr, 0), 0)
+            except socket.gaierror:
+                pass
+            else:
+                host.hostname = ni[0]
+        return host