Add support for multiple vendors

pybotvac/__init__.py

@@ -1,3 +1,5 @@
 from .account import Account
 from .robot import Robot
 from .version import __version__
+from .vorwerk import Vorwerk
+from .neato import Neato

pybotvac/account.py

@@ -22,14 +22,14 @@ class Account:
 
     """
 
-    ENDPOINT = 'https://beehive.neatocloud.com/'
-
-    def __init__(self, email, password):
+    def __init__(self, email, password, vendor):
         """Initialize the account data."""
         self._robots = set()
         self.robot_serials = {}
+        self._vendor = vendor
+        self._endpoint = vendor.endpoint
         self._maps = {}
-        self._headers = {'Accept': 'application/vnd.neato.nucleo.v1'}
+        self._headers = {'Accept': vendor.headers}
         self._login(email, password)
         self._persistent_maps = {}
 
@@ -41,7 +41,7 @@ def _login(self, email, password):
         :param password: Password for pybotvac account
         :return:
         """
-        response = requests.post(urljoin(self.ENDPOINT, 'sessions'),
+        response = requests.post(urljoin(self._endpoint, 'sessions'),
                                  json={'email': email,
                                        'password': password,
                                        'platform': 'ios',
@@ -84,7 +84,7 @@ def refresh_maps(self):
         """
         for robot in self.robots:
             resp2 = (
-                requests.get(urljoin(self.ENDPOINT, 'users/me/robots/{}/maps'.format(robot.serial)),
+                requests.get(urljoin(self._endpoint, 'users/me/robots/{}/maps'.format(robot.serial)),
                              headers=self._headers))
             resp2.raise_for_status()
             self._maps.update({robot.serial: resp2.json()})
@@ -95,7 +95,7 @@ def refresh_robots(self):
 
         :return:
         """
-        resp = requests.get(urljoin(self.ENDPOINT, 'dashboard'),
+        resp = requests.get(urljoin(self._endpoint, 'dashboard'),
                             headers=self._headers)
         resp.raise_for_status()
 
@@ -105,6 +105,7 @@ def refresh_robots(self):
 
             try:
                 self._robots.add(Robot(name=robot['name'],
+                                       vendor=self._vendor,
                                        serial=robot['serial'],
                                        secret=robot['secret_key'],
                                        traits=robot['traits'],
@@ -156,7 +157,7 @@ def refresh_persistent_maps(self):
         """
         for robot in self._robots:
             resp2 = (requests.get(urljoin(
-                self.ENDPOINT,
+                self._endpoint,
                 'users/me/robots/{}/persistent_maps'.format(robot.serial)),
                 headers=self._headers))
             resp2.raise_for_status()

pybotvac/neato.py

@@ -0,0 +1,8 @@
+import os
+
+class Neato:
+    name = 'neato'
+    endpoint = 'https://beehive.neatocloud.com/'
+    headers = 'application/vnd.neato.nucleo.v1'
+    cert_path = cert_path = os.path.join(os.path.dirname(__file__), 'cert',
+                                         'neatocloud.com.crt')

pybotvac/robot.py

@@ -18,7 +18,7 @@ class UnsupportedDevice(Exception):
 class Robot:
     """Data and methods for interacting with a Neato Botvac Connected vacuum robot"""
 
-    def __init__(self, serial, secret, traits, name='',
+    def __init__(self, serial, secret, traits, vendor, name='',
                  endpoint='https://nucleo.neatocloud.com:4443',
                  has_persistent_maps=False):
         """
@@ -30,13 +30,15 @@ def __init__(self, serial, secret, traits, name='',
         :param traits: Extras the robot supports
         """
         self.name = name
+        self._vendor = vendor
         self.serial = serial
         self.secret = secret
         self.traits = traits
         self.has_persistent_maps = has_persistent_maps
 
-        self._url = '{endpoint}/vendors/neato/robots/{serial}/messages'.format(
+        self._url = '{endpoint}/vendors/{vendor_name}/robots/{serial}/messages'.format(
             endpoint=re.sub(':\d+', '', endpoint),  # Remove port number
+            vendor_name=vendor.name,
             serial=self.serial)
         self._headers = {'Accept': 'application/vnd.neato.nucleo.v1'}
 
@@ -53,10 +55,9 @@ def _message(self, json):
         :return: server response
         """
 
-        cert_path = os.path.join(os.path.dirname(__file__), 'cert', 'neatocloud.com.crt')
         response = requests.post(self._url,
                                  json=json,
-                                 verify=cert_path,
+                                 verify=self._vendor.cert_path,
                                  auth=Auth(self.serial, self.secret),
                                  headers=self._headers)
         response.raise_for_status()

pybotvac/vorwerk.py

@@ -0,0 +1,7 @@
+class Vorwerk:
+    name = 'vorwerk'
+    endpoint = 'https://beehive.ksecosys.com/'
+    headers = 'application/vnd.neato.beehve.v1+json'
+    cert_path = False
+
+

sample/sample.py

@@ -1,14 +1,17 @@
 import sys
 
-from pybotvac import Account
+from pybotvac import Account, Vorwerk
 
 if sys.version_info[0] < 3:
     input = raw_input
 
 email = input('Enter email\n')
 password = input('Enter password\n')
 
-account = Account(email, password)
+vendor = Vorwerk()
+
+
+account = Account(email, password, vendor)
 
 print("Robots:")
 for robot in account.robots: