Merge pull request #113 from alfred82santa/master

Fixes & first steps in order to download encrypted files

.gitignore

@@ -10,4 +10,5 @@ chrome_cache/
 **/*.pyc
 profiles
 .idea
-docs/_build
+docs/_build
+.coverage

docs/conf.py

@@ -173,4 +173,4 @@
 # -- Options for todo extension ----------------------------------------------
 
 # If true, `todo` and `todoList` produce output, else they produce nothing.
-todo_include_todos = True
+todo_include_todos = True

requirements.txt

@@ -1,3 +1,5 @@
 python-dateutil>=2.6.0
 selenium>=3.4.3
 six>=1.10.0
+python-axolotl
+pycrypto

setup.py

@@ -9,13 +9,13 @@
 
 # To use a consistent encoding
 from codecs import open
-from os import path
+import os
 # Always prefer setuptools over distutils
 from setuptools import setup
 
 PACKAGE_NAME = 'webwhatsapi'
 
-path = path.join(path.dirname(__file__), PACKAGE_NAME, '__init__.py')
+path = os.path.join(os.path.dirname(__file__), PACKAGE_NAME, '__init__.py')
 
 with open(path, 'r') as file:
     t = compile(file.read(), path, 'exec', ast.PyCF_ONLY_AST)
@@ -43,7 +43,7 @@
             break
 
 # Get the long description from the README file
-with open(path.join(path.dirname(__file__), 'README.rst'), encoding='utf-8') as f:
+with open(os.path.join(os.path.dirname(__file__), 'README.rst'), encoding='utf-8') as f:
     long_description = f.read()
 
 setup(
@@ -95,11 +95,13 @@
 
     # You can just specify the packages manually here if your project is
     # simple. Or you can use find_packages().
-    packages=PACKAGE_NAME,
+    packages=[PACKAGE_NAME, ],
     install_requires=[
         'python-dateutil>=2.6.0',
         'selenium>=3.4.3',
         'six>=1.10.0',
+        'python-axolotl',
+        'pycrypto'
     ],
     extras_require={
     },

tox.ini

@@ -31,3 +31,4 @@ commands =
 [flake8]
 exclude = .tox,*.egg,build,data
 select = E,W,F
+max-line-length = 120

webwhatsapi/__init__.py

@@ -5,12 +5,18 @@
 
 """
 
+import binascii
 import logging
 from json import dumps, loads
 
 import os
 import shutil
 import tempfile
+from Crypto.Cipher import AES
+from axolotl.kdf.hkdfv3 import HKDFv3
+from axolotl.util.byteutil import ByteUtil
+from base64 import b64decode
+from io import BytesIO
 from selenium import webdriver
 from selenium.common.exceptions import NoSuchElementException
 from selenium.webdriver.common.by import By
@@ -19,10 +25,9 @@
 from selenium.webdriver.support import expected_conditions as EC
 from selenium.webdriver.support.ui import WebDriverWait
 
-from webwhatsapi.objects.chat import factory_chat, UserChat, Chat
-from webwhatsapi.objects.message import factory_message, MessageGroup
-from .consts import Selectors, URL
+from .objects.chat import UserChat, factory_chat
 from .objects.contact import Contact
+from .objects.message import MessageGroup, factory_message
 from .wapi_js_wrapper import WapiJsWrapper
 
 __version__ = '2.0.2'
@@ -141,10 +146,11 @@ def set_proxy(self, proxy):
         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, headless=False, autoconnect=True, logger=None):
+                 profile=None, headless=False, autoconnect=True, logger=None, extra_params=None):
         "Initialises the webdriver"
 
         self.logger = logger or self.logger
+        extra_params = extra_params or {}
 
         if profile is not None:
             self._profile_path = profile
@@ -161,7 +167,7 @@ def __init__(self, client="firefox", username="API", proxy=None, command_executo
                 self._profile = webdriver.FirefoxProfile(self._profile_path)
             else:
                 self._profile = webdriver.FirefoxProfile()
-            if loadstyles == False:
+            if not loadstyles:
                 # Disable CSS
                 self._profile.set_preference('permissions.default.stylesheet', 2)
                 # Disable images
@@ -183,21 +189,22 @@ def __init__(self, client="firefox", username="API", proxy=None, command_executo
             capabilities['webStorageEnabled'] = True
 
             self.logger.info("Starting webdriver")
-            self.driver = webdriver.Firefox(capabilities=capabilities, options=options)
+            self.driver = webdriver.Firefox(capabilities=capabilities, options=options, **extra_params)
 
         elif self.client == "chrome":
             self._profile = webdriver.chrome.options.Options()
             if self._profile_path is not None:
                 self._profile.add_argument("user-data-dir=%s" % self._profile_path)
             if proxy is not None:
                 profile.add_argument('--proxy-server=%s' % proxy)
-            self.driver = webdriver.Chrome(chrome_options=self._profile)
+            self.driver = webdriver.Chrome(chrome_options=self._profile, **extra_params)
 
         elif client == 'remote':
             capabilities = DesiredCapabilities.FIREFOX.copy()
             self.driver = webdriver.Remote(
                 command_executor=command_executor,
-                desired_capabilities=capabilities
+                desired_capabilities=capabilities,
+                **extra_params
             )
 
         else:
@@ -395,5 +402,32 @@ def group_get_admins(self, group_id):
         for admin_id in admin_ids:
             yield self.get_contact_from_id(admin_id)
 
+    def download_file(self, url):
+        return b64decode(self.wapi_functions.downloadFile(url))
+
+    def download_media(self, media_msg):
+        try:
+            if media_msg.content:
+                return BytesIO(b64decode(self.content))
+        except AttributeError:
+            pass
+
+        file_data = self.download_file(media_msg.client_url)
+
+        media_key = b64decode(media_msg.media_key)
+        derivative = HKDFv3().deriveSecrets(media_key,
+                                            binascii.unhexlify(media_msg.crypt_keys[media_msg.type]),
+                                            112)
+
+        parts = ByteUtil.split(derivative, 16, 32)
+        iv = parts[0]
+        cipher_key = parts[1]
+        e_file = file_data[:-10]
+
+        AES.key_size = 128
+        cr_obj = AES.new(key=cipher_key, mode=AES.MODE_CBC, IV=iv)
+
+        return BytesIO(cr_obj.decrypt(e_file))
+
     def quit(self):
         self.driver.quit()

webwhatsapi/async_driver.py

@@ -1,7 +1,15 @@
 from asyncio import get_event_loop
+
+import binascii
+
+from Crypto.Cipher import AES
+from axolotl.kdf.hkdfv3 import HKDFv3
+from axolotl.util.byteutil import ByteUtil
+from base64 import b64decode
 from concurrent.futures import ThreadPoolExecutor
 
 from functools import partial
+from io import BytesIO
 
 from webwhatsapi import factory_message
 from . import WhatsAPIDriver
@@ -10,14 +18,14 @@
 class WhatsAPIDriverAsync:
 
     def __init__(self, client="firefox", username="API", proxy=None, command_executor=None, loadstyles=False,
-                 profile=None, headless=False, logger=None, loop=None):
+                 profile=None, headless=False, logger=None, extra_params=None, loop=None):
 
         self._driver = WhatsAPIDriver(client=client, username=username, proxy=proxy, command_executor=command_executor,
                                       loadstyles=loadstyles, profile=profile, headless=headless, logger=logger,
-                                      autoconnect=False)
+                                      autoconnect=False, extra_params=extra_params)
 
         self.loop = loop or get_event_loop()
-        self._pool_executor = ThreadPoolExecutor(max_workers=4)
+        self._pool_executor = ThreadPoolExecutor(max_workers=1)
 
     async def get_local_storage(self):
         return await self.loop.run_in_executor(self._pool_executor, self._driver.get_local_storage)
@@ -130,5 +138,35 @@ async def group_get_admins(self, group_id):
         for admin_id in admin_ids:
             yield await self.get_contact_from_id(admin_id)
 
+    async def download_file(self, url):
+        return await self.loop.run_in_executor(self._pool_executor,
+                                               self._driver.download_file,
+                                               url)
+
+    async def download_media(self, media_msg):
+        try:
+            if media_msg.content:
+                return BytesIO(b64decode(self.content))
+        except AttributeError:
+            pass
+
+        file_data = await self.download_file(media_msg.client_url)
+
+        media_key = b64decode(media_msg.media_key)
+        derivative = HKDFv3().deriveSecrets(media_key,
+                                            binascii.unhexlify(media_msg.crypt_keys[media_msg.type]),
+                                            112)
+
+        parts = ByteUtil.split(derivative, 16, 32)
+        iv = parts[0]
+        cipher_key = parts[1]
+        e_file = file_data[:-10]
+
+        AES.key_size = 128
+        cr_obj = AES.new(key=cipher_key, mode=AES.MODE_CBC, IV=iv)
+
+        return BytesIO(cr_obj.decrypt(e_file))
+
     async def quit(self):
         return await self.loop.run_in_executor(self._pool_executor, self._driver.quit)
+

webwhatsapi/js/wapi.js

@@ -60,21 +60,41 @@ window.WAPI._serializeNotificationObj = (obj) => ({
 });
 
 //TODO: Add chat ref
-window.WAPI._serializeMessageObj = (obj) => ({
-    sender: WAPI._serializeContactObj(obj["senderObj"]),
-    timestamp: obj["t"],
-    content: obj["body"],
-    isGroupMsg: obj.__x_isGroupMsg,
-    isLink: obj.__x_isLink,
-    isMMS: obj.__x_isMMS,
-    isMedia: obj.__x_isMedia,
-    isNotification: obj.__x_isNotification,
-    isPSA: obj.__x_isPSA,
-    type: obj.__x_type,
-    size: obj.__x_size,
-    mime: obj.__x_mimetype,
-});
+window.WAPI._serializeMessageObj = function(obj) {
+
+    let data = {
+        sender: WAPI._serializeContactObj(obj["senderObj"]),
+        id: obj.id._serialized,
+        timestamp: obj["t"],
+        content: obj["body"],
+        isGroupMsg: obj.__x_isGroupMsg,
+        isLink: obj.__x_isLink,
+        isMMS: obj.__x_isMMS,
+        isMedia: obj.__x_isMedia,
+        isNotification: obj.__x_isNotification,
+        isPSA: obj.__x_isPSA,
+        type: obj.__x_type,
+        size: obj.__x_size,
+        mime: obj.__x_mimetype,
+        chatId: obj.__x_id.remote
+    }
 
+    if (data.isMedia || data.isMMS) {
+        data['clientUrl'] = obj['__x_clientUrl'];
+        data['mediaKey'] = obj['__x_mediaKey'];
+        data['mediaData'] = {
+            duration: obj['__x_mediaData']['__x_duration'],
+            filehash: obj['__x_mediaData']['__x_filehash'],
+            mimetype: obj['__x_mediaData']['__x_mimetype'],
+            encriptationKey: obj['__x_mediaData']['__x_encryptionKey'],
+            fullHeight: obj['__x_mediaData']['__x_fullHeight'],
+            fullWidth: obj['__x_mediaData']['__x_fullWidth'],
+            size: obj['__x_mediaData']['__x_size'],
+        }
+    }
+
+    return data
+}
 /**
  * Fetches all contact objects from store
  *
@@ -444,4 +464,24 @@ window.WAPI.getCommonGroups = async function (id, done) {
     return output;
 };
 
+window.WAPI.downloadFile = function (url, done) {
+    let xhr = new XMLHttpRequest();
+
+    xhr.onload = function() {
+        if (xhr.readyState == 4) {
+            if (xhr.status == 200) {
+                let reader = new FileReader();
+                reader.readAsDataURL(xhr.response);
+                reader.onload =  function(e){
+                    done(reader.result.substr(reader.result.indexOf(',')+1))
+                };
+            } else {
+                console.error(xhr.statusText);
+            }
+        }
+    };
+    xhr.open("GET", url, true);
+    xhr.responseType = 'blob';
+    xhr.send(null);
+}
 

webwhatsapi/objects/contact.py

@@ -8,6 +8,7 @@ class Contact(WhatsappObjectWithId):
     """
     Class which represents a Contact on user's phone
     """
+
     def __init__(self, js_obj, driver=None):
         """
 

webwhatsapi/objects/message.py

@@ -35,8 +35,12 @@ def __init__(self, js_obj, driver=None):
         :type js_obj: dict
         """
         super(Message, self).__init__(js_obj, driver)
-        self.sender = False if js_obj["sender"] == False else Contact(js_obj["sender"], driver)
+
+        self.id = js_obj["id"]
+        self.sender = False if js_obj["sender"] is False else Contact(js_obj["sender"], driver)
         self.timestamp = datetime.fromtimestamp(js_obj["timestamp"])
+        self.chat_id = js_obj['chatId']
+
         if js_obj["content"]:
             self.content = js_obj["content"]
             self.safe_content = safe_str(self.content[0:25]) + '...'
@@ -50,18 +54,26 @@ def __repr__(self):
 
 
 class MediaMessage(Message):
+    crypt_keys = {'document': '576861747341707020446f63756d656e74204b657973',
+                  'image': '576861747341707020496d616765204b657973',
+                  'video': '576861747341707020566964656f204b657973',
+                  'ptt': '576861747341707020417564696f204b657973'}
+
     def __init__(self, js_obj, driver=None):
         super(MediaMessage, self).__init__(js_obj, driver)
 
         self.type = self.js_obj["type"]
         self.size = self.js_obj["size"]
         self.mime = self.js_obj["mime"]
 
+        self.media_key = self.js_obj.get('mediaKey')
+        self.client_url = self.js_obj.get('clientUrl')
+
         extension = mimetypes.guess_extension(self.mime)
         try:
-            self.filename = ''.join([self.js_obj["__x_filehash"], extension])
+            self.filename = ''.join([self.js_obj["filehash"], extension])
         except KeyError:
-            self.filename = ''.join([str(id(self)), extension])
+            self.filename = ''.join([str(id(self)), extension or ''])
 
     def save_media(self, path):
         with open(os.path.join(path, self.filename), "wb") as output: