From 38030fcfca91d067d8783579150dcc3ecb4d283f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Per=20Sandstr=C3=B6m?= <per.j.sandstrom@gmail.com>
Date: Wed, 15 Jun 2016 07:17:32 +0200
Subject: [PATCH] ASUSWRT Autodetect protocol (#2300)

---
 .../components/device_tracker/asuswrt.py      | 60 +++++++++++++------
 1 file changed, 42 insertions(+), 18 deletions(-)

diff --git a/homeassistant/components/device_tracker/asuswrt.py b/homeassistant/components/device_tracker/asuswrt.py
index e049ed87c48..282ae46ba85 100644
--- a/homeassistant/components/device_tracker/asuswrt.py
+++ b/homeassistant/components/device_tracker/asuswrt.py
@@ -6,8 +6,10 @@ https://home-assistant.io/components/device_tracker.asuswrt/
 """
 import logging
 import re
+import socket
 import telnetlib
 import threading
+from collections import namedtuple
 from datetime import timedelta
 
 from homeassistant.components.device_tracker import DOMAIN
@@ -68,6 +70,8 @@ def get_scanner(hass, config):
 
     return scanner if scanner.success_init else None
 
+AsusWrtResult = namedtuple('AsusWrtResult', 'neighbors leases arp')
+
 
 class AsusWrtDeviceScanner(object):
     """This class queries a router running ASUSWRT firmware."""
@@ -130,7 +134,7 @@ class AsusWrtDeviceScanner(object):
 
     def ssh_connection(self):
         """Retrieve data from ASUSWRT via the ssh protocol."""
-        from pexpect import pxssh
+        from pexpect import pxssh, exceptions
 
         try:
             ssh = pxssh.pxssh()
@@ -140,7 +144,7 @@ class AsusWrtDeviceScanner(object):
                 ssh.login(self.host, self.username, self.password)
             else:
                 _LOGGER.error('No password or public key specified')
-                return('', '', '')
+                return None
             ssh.sendline(_IP_NEIGH_CMD)
             ssh.prompt()
             neighbors = ssh.before.split(b'\n')[1:-1]
@@ -157,10 +161,13 @@ class AsusWrtDeviceScanner(object):
                 ssh.prompt()
                 leases_result = ssh.before.split(b'\n')[1:-1]
             ssh.logout()
-            return (neighbors, leases_result, arp_result)
+            return AsusWrtResult(neighbors, leases_result, arp_result)
         except pxssh.ExceptionPxssh as exc:
-            _LOGGER.exception('Unexpected response from router: %s', exc)
-            return ('', '', '')
+            _LOGGER.error('Unexpected response from router: %s', exc)
+            return None
+        except exceptions.EOF:
+            _LOGGER.error('Connection refused or no route to host')
+            return None
 
     def telnet_connection(self):
         """Retrieve data from ASUSWRT via the telnet protocol."""
@@ -186,25 +193,42 @@ class AsusWrtDeviceScanner(object):
                 leases_result = (telnet.read_until(prompt_string).
                                  split(b'\n')[1:-1])
             telnet.write('exit\n'.encode('ascii'))
-            return (neighbors, leases_result, arp_result)
+            return AsusWrtResult(neighbors, leases_result, arp_result)
         except EOFError:
-            _LOGGER.exception("Unexpected response from router")
-            return ('', '', '')
+            _LOGGER.error("Unexpected response from router")
+            return None
         except ConnectionRefusedError:
-            _LOGGER.exception("Connection refused by router,"
-                              " is telnet enabled?")
-            return ('', '', '')
+            _LOGGER.error("Connection refused by router, is telnet enabled?")
+            return None
+        except socket.gaierror as exc:
+            _LOGGER.error("Socket exception: %s", exc)
+            return None
+        except OSError as exc:
+            _LOGGER.error("OSError: %s", exc)
+            return None
 
     def get_asuswrt_data(self):
         """Retrieve data from ASUSWRT and return parsed result."""
-        if self.protocol == 'telnet':
-            neighbors, leases_result, arp_result = self.telnet_connection()
+        if self.protocol == 'ssh':
+            result = self.ssh_connection()
+        elif self.protocol == 'telnet':
+            result = self.telnet_connection()
         else:
-            neighbors, leases_result, arp_result = self.ssh_connection()
+            # autodetect protocol
+            result = self.ssh_connection()
+            if result:
+                self.protocol = 'ssh'
+            else:
+                result = self.telnet_connection()
+                if result:
+                    self.protocol = 'telnet'
+
+        if not result:
+            return {}
 
         devices = {}
         if self.mode == 'ap':
-            for lease in leases_result:
+            for lease in result.leases:
                 match = _WL_REGEX.search(lease.decode('utf-8'))
 
                 if not match:
@@ -214,7 +238,7 @@ class AsusWrtDeviceScanner(object):
                 host = ''
 
                 # match mac addresses to IP addresses in ARP table
-                for arp in arp_result:
+                for arp in result.arp:
                     if match.group('mac').lower() in arp.decode('utf-8'):
                         arp_match = _ARP_REGEX.search(arp.decode('utf-8'))
                         if not arp_match:
@@ -228,7 +252,7 @@ class AsusWrtDeviceScanner(object):
                             'mac': match.group('mac').upper(),
                             }
         else:
-            for lease in leases_result:
+            for lease in result.leases:
                 match = _LEASES_REGEX.search(lease.decode('utf-8'))
 
                 if not match:
@@ -248,7 +272,7 @@ class AsusWrtDeviceScanner(object):
                     'mac': match.group('mac').upper(),
                     }
 
-        for neighbor in neighbors:
+        for neighbor in result.neighbors:
             match = _IP_NEIGH_REGEX.search(neighbor.decode('utf-8'))
             if not match:
                 _LOGGER.warning("Could not parse neighbor row: %s", neighbor)