Resolved problems with Python3; Removed prints and exits; Fix save firefox profile; Store WebStorage in firefox profile in order to keep session; Use local storage if it exists in profile

alfred82santa · alfred82santa · commit 2c010ab20d9e · 2018-02-15T18:28:54.000+01:00
diff --git a/.env b/.env
@@ -0,0 +1 @@
+webwhatsapi@2.0.1
diff --git a/.gitignore b/.gitignore
@@ -8,3 +8,5 @@ chrome_cache/
 *.png
 **/.DS_Store
 **/*.pyc
+profiles
+.idea
diff --git a/setup.py b/setup.py
@@ -5,16 +5,45 @@
 https://github.com/pypa/sampleproject
 """
 
-# Always prefer setuptools over distutils
-from setuptools import setup, find_packages
+import ast
+
 # To use a consistent encoding
 from codecs import open
 from os import path
-
-here = path.abspath(path.dirname(__file__))
+# Always prefer setuptools over distutils
+from setuptools import setup
+
+PACKAGE_NAME = 'webwhatsapi'
+
+path = path.join(path.dirname(__file__), PACKAGE_NAME, '__init__.py')
+
+with open(path, 'r') as file:
+    t = compile(file.read(), path, 'exec', ast.PyCF_ONLY_AST)
+    for node in (n for n in t.body if isinstance(n, ast.Assign)):
+        if len(node.targets) != 1:
+            continue
+
+        name = node.targets[0]
+        if not isinstance(name, ast.Name) or \
+                name.id not in ('__version__', '__version_info__', 'VERSION'):
+            continue
+
+        v = node.value
+        if isinstance(v, ast.Str):
+            version = v.s
+            break
+        if isinstance(v, ast.Tuple):
+            r = []
+            for e in v.elts:
+                if isinstance(e, ast.Str):
+                    r.append(e.s)
+                elif isinstance(e, ast.Num):
+                    r.append(str(e.n))
+            version = '.'.join(r)
+            break
 
 # Get the long description from the README file
-with open(path.join(here, 'README.rst'), encoding='utf-8') as f:
+with open(path.join(path.dirname(__file__), 'README.rst'), encoding='utf-8') as f:
     long_description = f.read()
 
 setup(
@@ -23,14 +52,14 @@
     # Versions should comply with PEP440.  For a discussion on single-sourcing
     # the version across setup.py and the project code, see
     # https://packaging.python.org/en/latest/single_source_version.html
-    version='2.0.1',
+    version=version,
 
     description='A python interface for Whatsapp Web',
     long_description=long_description,
 
     # The project's main homepage.
     url='https://github.com/mukulhase/WhatsAPI',
-    download_url='https://github.com/mukulhase/WhatsAPI/archive/2.0.1.tar.gz',
+    download_url='https://github.com/mukulhase/WhatsAPI/archive/{}.tar.gz'.format(version),
 
     # Author details
     author='Mukul Hase',
@@ -58,14 +87,15 @@
         # Specify the Python versions you support here. In particular, ensure
         # that you indicate whether you support Python 2, Python 3 or both.
         'Programming Language :: Python :: 2.7',
+        'Programming Language :: Python :: 3.6',
     ],
 
     # What does your project relate to?
     keywords='Whatsapp Chat Bot Chatbot Selenium Web Whatsapp API',
 
     # You can just specify the packages manually here if your project is
     # simple. Or you can use find_packages().
-    packages=find_packages(exclude=['contrib', 'docs', 'tests']),
+    packages=PACKAGE_NAME,
     install_requires=[
         'python-dateutil>=2.6.0',
         'selenium>=3.4.3',
diff --git a/webwhatsapi/__init__.py b/webwhatsapi/__init__.py
@@ -2,25 +2,27 @@
 WhatsAPI module
 """
 
+import logging
+from json import dumps, loads
 
-from __future__ import print_function
-
-import sys
-import datetime
-import time
 import os
-import sys
-import logging
-import pickle
 import tempfile
-
+import shutil
 from selenium import webdriver
-from selenium.webdriver.chrome.options import Options
-from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
+from selenium.common.exceptions import NoSuchElementException
 from selenium.webdriver.common.by import By
+from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
+from selenium.webdriver.firefox.options import Options
 from selenium.webdriver.support import expected_conditions as EC
 from selenium.webdriver.support.ui import WebDriverWait
-from selenium.common.exceptions import NoSuchElementException
+
+from .consts import Selectors, URL
+from .objects.chat import Chat, GroupChat, UserChat
+from .objects.contact import Contact
+from .objects.message import Message, MessageGroup
+from .wapi_js_wrapper import WapiJsWrapper
+
+__version__ = '2.0.2'
 
 
 class WhatsAPIDriverStatus(object):
@@ -30,21 +32,18 @@ class WhatsAPIDriverStatus(object):
     NotLoggedIn = 'NotLoggedIn'
     LoggedIn = 'LoggedIn'
 
-from consts import Selectors, URL
-from objects.chat import Chat, GroupChat, UserChat
-from objects.contact import Contact
-from objects.message import Message, MessageGroup
-from webwhatsapi.wapi_js_wrapper import WapiJsWrapper
 
-if __debug__:
-    import pprint
-    pp = pprint.PrettyPrinter(indent=4)
+class WhatsAPIException(Exception):
+    pass
+
 
 class WhatsAPIDriver(object):
     _PROXY = None
 
     _URL = "https://web.whatsapp.com"
 
+    _LOCAL_STORAGE_FILE = 'localStorage.json'
+
     _SELECTORS = {
         'firstrun': "#wrapper",
         'qrCode': "img[alt=\"Scan me!\"]",
@@ -78,13 +77,40 @@ class WhatsAPIDriver(object):
     # Do not alter this
     _profile = None
 
-    def save_firefox_profile(self):
+    def get_local_storage(self):
+        return self.driver.execute_script('return window.localStorage;')
+
+    def set_local_storage(self, data):
+        self.driver.execute_script(''.join(["window.localStorage.setItem('{}', '{}');".format(k, v)
+                                            for k, v in data.items()]))
+
+    def save_firefox_profile(self, remove_old=False):
         "Function to save the firefox profile to the permanant one"
         self.logger.info("Saving profile from %s to %s" % (self._profile.path, self._profile_path))
-        os.system("cp -R " + self._profile.path + " "+ self._profile_path)
-        cookie_file = os.path.join(self._profile_path, "cookies.pkl")
-        if self.driver:
-            pickle.dump(self.driver.get_cookies() , open(cookie_file,"wb"))
+
+        if remove_old:
+            if os.path.exists(self._profile_path):
+                try:
+                    shutil.rmtree(self._profile_path)
+                except OSError:
+                    pass
+
+            shutil.copytree(os.path.join(self._profile.path), self._profile_path,
+                            ignore=shutil.ignore_patterns("parent.lock", "lock", ".parentlock"))
+        else:
+            for item in os.listdir(self._profile.path):
+                if item in ["parent.lock", "lock", ".parentlock"]:
+                    continue
+                s = os.path.join(self._profile.path, item)
+                d = os.path.join(self._profile_path, item)
+                if os.path.isdir(s):
+                    shutil.copytree(s, d,
+                                    ignore=shutil.ignore_patterns("parent.lock", "lock", ".parentlock"))
+                else:
+                    shutil.copy2(s, d)
+
+        with open(os.path.join(self._profile_path, self._LOCAL_STORAGE_FILE), 'w') as f:
+            f.write(dumps(self.get_local_storage()))
 
     def set_proxy(self, proxy):
         self.logger.info("Setting proxy to %s" % proxy)
@@ -95,7 +121,8 @@ def set_proxy(self, proxy):
         self._profile.set_preference("network.proxy.ssl", proxy_address)
         self._profile.set_preference("network.proxy.ssl_port", int(proxy_port))
 
-    def __init__(self, client="firefox", username="API", proxy=None, command_executor=None, loadstyles=False, profile=None):
+    def __init__(self, client="firefox", username="API", proxy=None, command_executor=None, loadstyles=False,
+                 profile=None, headless=False):
         "Initialises the webdriver"
 
         # Get the name of the config folder
@@ -105,8 +132,8 @@ def __init__(self, client="firefox", username="API", proxy=None, command_executo
             if not os.path.exists(self.config_dir):
                 os.makedirs(self.config_dir)
         except OSError:
-            print("Error: Could not create config dir")
-            exit(-1)
+            self.logger.critical("Error: Could not create config dir")
+            raise WhatsAPIException("Error: Could not create config dir")
 
         self.logger.setLevel(logging.DEBUG)
 
@@ -119,9 +146,8 @@ def __init__(self, client="firefox", username="API", proxy=None, command_executo
             self._profile_path = profile
             self.logger.info("Checking for profile at %s" % self._profile_path)
             if not os.path.exists(self._profile_path):
-                print("Could not find profile at %s" % profile)
-                self.logger.error("Could not find profile at %s" % profile)
-                exit(-1)
+                self.logger.critical("Could not find profile at %s" % profile)
+                raise WhatsAPIException("Could not find profile at %s" % profile)
         else:
             self._profile_path = None
 
@@ -138,11 +164,22 @@ def __init__(self, client="firefox", username="API", proxy=None, command_executo
                 self._profile.set_preference('permissions.default.image', 2)
                 ## Disable Flash
                 self._profile.set_preference('dom.ipc.plugins.enabled.libflashplayer.so',
-                                              'false')
+                                             'false')
             if proxy is not None:
                 self.set_proxy(proxy)
+
+            options = Options()
+
+            if headless:
+                options.set_headless()
+
+            options.profile = self._profile
+
+            capabilities = DesiredCapabilities.FIREFOX.copy()
+            capabilities['webStorageEnabled'] = True
+
             self.logger.info("Starting webdriver")
-            self.driver = webdriver.Firefox(self._profile)
+            self.driver = webdriver.Firefox(capabilities=capabilities, options=options)
 
         elif self.client == "chrome":
             self._profile = webdriver.chrome.options.Options()
@@ -161,14 +198,20 @@ def __init__(self, client="firefox", username="API", proxy=None, command_executo
 
         else:
             self.logger.error("Invalid client: %s" % client)
-            print("Enter a valid client name")
         self.username = username
         self.wapi_functions = WapiJsWrapper(self.driver)
 
         self.driver.set_script_timeout(500)
         self.driver.implicitly_wait(10)
         self.driver.get(self._URL)
 
+        local_storage_file = os.path.join(self._profile.path, self._LOCAL_STORAGE_FILE)
+        if os.path.exists(local_storage_file):
+            with open(local_storage_file) as f:
+                self.set_local_storage(loads(f.read()))
+
+            self.driver.refresh()
+
     def wait_for_login(self):
         """Waits for the QR to go away"""
         WebDriverWait(self.driver, 90).until(
@@ -182,7 +225,6 @@ def get_qr(self):
         qr = self.driver.find_element_by_css_selector(self._SELECTORS['qrCode'])
         fd, fn_png = tempfile.mkstemp(prefix=self.username, suffix='.png')
         self.logger.debug("QRcode image saved at %s" % fn_png)
-        print(fn_png)
         qr.screenshot(fn_png)
         os.close(fd)
         return fn_png
@@ -288,7 +330,7 @@ def get_chat_from_phone_number(self, number):
         :return: Chat
         :rtype: Chat
         """
-        chats = filter(lambda chat:(type(chat) is UserChat), self.get_all_chats())
+        chats = filter(lambda chat: (type(chat) is UserChat), self.get_all_chats())
         return next((contact for contact in chats if (number in contact.id)), None)
 
     def reload_qr(self):
diff --git a/webwhatsapi/objects/chat.py b/webwhatsapi/objects/chat.py
@@ -1,8 +1,10 @@
-from webwhatsapi.objects.whatsapp_object import WhatsappObject, driver_needed
-from webwhatsapi.helper import safe_str
-from webwhatsapi.objects.message import Message
+from six import with_metaclass
+
+from .message import Message
+from .whatsapp_object import WhatsappObject, driver_needed
+from ..helper import safe_str
+
 
-##TODO: Fix relative imports for Python3
 class ChatMetaClass(type):
     """
     Chat type factory
@@ -28,8 +30,7 @@ def __call__(cls, js_obj, driver=None):
         return type.__call__(UserChat, js_obj, driver)
 
 
-class Chat(WhatsappObject):
-    __metaclass__ = ChatMetaClass
+class Chat(with_metaclass(ChatMetaClass, WhatsappObject)):
 
     def __init__(self, js_obj, driver=None):
         super(Chat, self).__init__(js_obj, driver)
@@ -52,6 +53,7 @@ def load_earlier_messages(self):
     def load_all_earlier_messages(self):
         self.driver.wapi_functions.loadAllEarlierMessages(self.id)
 
+
 class UserChat(Chat):
     def __init__(self, js_obj, driver=None):
         super(UserChat, self).__init__(js_obj, driver)
diff --git a/webwhatsapi/objects/contact.py b/webwhatsapi/objects/contact.py
@@ -1,6 +1,8 @@
-from webwhatsapi.objects.whatsapp_object import WhatsappObject, driver_needed
-from webwhatsapi.helper import safe_str
-import six
+from six import string_types
+
+from .whatsapp_object import WhatsappObject, driver_needed
+from ..helper import safe_str
+
 
 class Contact(WhatsappObject):
     def __init__(self, js_obj, driver=None):
@@ -19,7 +21,7 @@ def get_chat(self):
 
     def get_safe_name(self):
         name = (self.name or self.push_name or self.formatted_name)
-        if (isinstance(name, six.string_types)):
+        if (isinstance(name, string_types)):
             safe_name = safe_str(name)
         else:
             safe_name = "Unknown"
diff --git a/webwhatsapi/objects/message.py b/webwhatsapi/objects/message.py
diff --git a/webwhatsapi/wapi_js_wrapper.py b/webwhatsapi/wapi_js_wrapper.py
