Allow Twitter notifications to include media

homeassistant/components/notify/twitter.py

@@ -4,13 +4,16 @@
 For more details about this platform, please refer to the documentation at
 https://home-assistant.io/components/notify.twitter/
 """
+import json
 import logging
+import mimetypes
+import os
 
 import voluptuous as vol
 
 import homeassistant.helpers.config_validation as cv
 from homeassistant.components.notify import (
-    PLATFORM_SCHEMA, BaseNotificationService)
+    ATTR_DATA, PLATFORM_SCHEMA, BaseNotificationService)
 from homeassistant.const import CONF_ACCESS_TOKEN, CONF_USERNAME
 
 REQUIREMENTS = ['TwitterAPI==2.4.5']
@@ -21,6 +24,8 @@
 CONF_CONSUMER_SECRET = 'consumer_secret'
 CONF_ACCESS_TOKEN_SECRET = 'access_token_secret'
 
+ATTR_MEDIA = 'media'
+
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
     vol.Required(CONF_ACCESS_TOKEN): cv.string,
     vol.Required(CONF_ACCESS_TOKEN_SECRET): cv.string,
@@ -33,6 +38,7 @@
 def get_service(hass, config, discovery_info=None):
     """Get the Twitter notification service."""
     return TwitterNotificationService(
+        hass,
         config[CONF_CONSUMER_KEY], config[CONF_CONSUMER_SECRET],
         config[CONF_ACCESS_TOKEN], config[CONF_ACCESS_TOKEN_SECRET],
         config.get(CONF_USERNAME)
@@ -42,26 +48,113 @@ def get_service(hass, config, discovery_info=None):
 class TwitterNotificationService(BaseNotificationService):
     """Implementation of a notification service for the Twitter service."""
 
-    def __init__(self, consumer_key, consumer_secret, access_token_key,
+    def __init__(self, hass, consumer_key, consumer_secret, access_token_key,
                  access_token_secret, username):
         """Initialize the service."""
         from TwitterAPI import TwitterAPI
         self.user = username
+        self.hass = hass
         self.api = TwitterAPI(consumer_key, consumer_secret, access_token_key,
                               access_token_secret)
 
     def send_message(self, message="", **kwargs):
-        """Tweet a message."""
+        """Tweet a message, optionally with media."""
+        data = kwargs.get(ATTR_DATA)
+        media = data.get(ATTR_MEDIA)
+        if not self.hass.config.is_allowed_path(media):
+            _LOGGER.warning("'%s' is not in a whitelisted area.", media)
+            return
+
+        media_id = self.upload_media(media)
+
         if self.user:
-            resp = self.api.request(
-                'direct_messages/new', {'text': message, 'user': self.user})
+            resp = self.api.request('direct_messages/new',
+                                    {'text': message, 'user': self.user,
+                                     'media_ids': media_id})
         else:
-            resp = self.api.request('statuses/update', {'status': message})
+            resp = self.api.request('statuses/update',
+                                    {'status': message, 'media_ids': media_id})
 
         if resp.status_code != 200:
-            import json
-            obj = json.loads(resp.text)
-            error_message = obj['errors'][0]['message']
-            error_code = obj['errors'][0]['code']
-            _LOGGER.error("Error %s : %s (Code %s)", resp.status_code,
-                          error_message, error_code)
+            self.log_error_resp(resp)
+
+    def upload_media(self, media_path=None):
+        """Upload media."""
+        if not media_path:
+            return None
+
+        (media_type, _) = mimetypes.guess_type(media_path)
+        total_bytes = os.path.getsize(media_path)
+
+        file = open(media_path, 'rb')
+        resp = self.upload_media_init(media_type, total_bytes)
+
+        if 199 > resp.status_code < 300:
+            self.log_error_resp(resp)
+            return None
+
+        media_id = resp.json()['media_id']
+        media_id = self.upload_media_chunked(file, total_bytes,
+                                             media_id)
+
+        resp = self.upload_media_finalize(media_id)
+        if 199 > resp.status_code < 300:
+            self.log_error_resp(resp)
+
+        return media_id
+
+    def upload_media_init(self, media_type, total_bytes):
+        """Upload media, INIT phase."""
+        resp = self.api.request('media/upload',
+                                {'command': 'INIT', 'media_type': media_type,
+                                 'total_bytes': total_bytes})
+        return resp
+
+    def upload_media_chunked(self, file, total_bytes, media_id):
+        """Upload media, chunked append."""
+        segment_id = 0
+        bytes_sent = 0
+        while bytes_sent < total_bytes:
+            chunk = file.read(4 * 1024 * 1024)
+            resp = self.upload_media_append(chunk, media_id, segment_id)
+            if resp.status_code not in range(200, 299):
+                self.log_error_resp_append(resp)
+                return None
+            segment_id = segment_id + 1
+            bytes_sent = file.tell()
+            self.log_bytes_sent(bytes_sent, total_bytes)
+        return media_id
+
+    def upload_media_append(self, chunk, media_id, segment_id):
+        """Upload media, append phase."""
+        return self.api.request('media/upload',
+                                {'command': 'APPEND', 'media_id': media_id,
+                                 'segment_index': segment_id},
+                                {'media': chunk})
+
+    def upload_media_finalize(self, media_id):
+        """Upload media, finalize phase."""
+        return self.api.request('media/upload',
+                                {'command': 'FINALIZE', 'media_id': media_id})
+
+    @staticmethod
+    def log_bytes_sent(bytes_sent, total_bytes):
+        """Log upload progress."""
+        _LOGGER.debug("%s of %s bytes uploaded", str(bytes_sent),
+                      str(total_bytes))
+
+    @staticmethod
+    def log_error_resp(resp):
+        """Log error response."""
+        obj = json.loads(resp.text)
+        error_message = obj['error']
+        _LOGGER.error("Error %s : %s", resp.status_code, error_message)
+
+    @staticmethod
+    def log_error_resp_append(resp):
+        """Log error response, during upload append phase."""
+        obj = json.loads(resp.text)
+        error_message = obj['errors'][0]['message']
+        error_code = obj['errors'][0]['code']
+        _LOGGER.error("Error %s : %s (Code %s)", resp.status_code,
+                      error_message, error_code)