Pulled new API from PUTIO

libs/pio/api.py

@@ -1,26 +1,53 @@
 # -*- coding: utf-8 -*-
-
 # Changed 
 # Removed iso8601 library requirement 
 # Added CP logging
 
 import os
-import re
 import json
+import binascii
 import webbrowser
+try:
     from urllib import urlencode
-from couchpotato import CPLog
+except ImportError:
-from dateutil.parser import parse
+    from urllib.parse import urlencode
+from datetime import datetime
 
+import tus
 import requests
+from requests.adapters import HTTPAdapter
+from requests.packages.urllib3.util.retry import Retry
+from couchpotato import CPLog
+
+
+KB = 1024
+MB = 1024 * KB
+
+# Read and write operations are limited to this chunk size.
+# This can make a big difference when dealing with large files.
+CHUNK_SIZE = 256 * KB
 
 BASE_URL = 'https://api.put.io/v2'
+UPLOAD_URL = 'https://upload.put.io/v2/files/upload'
+TUS_UPLOAD_URL = 'https://upload.put.io/files/'
 ACCESS_TOKEN_URL = 'https://api.put.io/v2/oauth2/access_token'
 AUTHENTICATION_URL = 'https://api.put.io/v2/oauth2/authenticate'
 
 log = CPLog(__name__)
 
 
+class APIError(Exception):
+    pass
+
+
+class ClientError(APIError):
+    pass
+
+
+class ServerError(APIError):
+    pass
+
+
 class AuthHelper(object):
 
     def __init__(self, client_id, client_secret, redirect_uri, type='code'):
@@ -58,10 +85,21 @@ class AuthHelper(object):
 
 class Client(object):
 
-    def __init__(self, access_token):
+    def __init__(self, access_token, use_retry=False):
         self.access_token = access_token
         self.session = requests.session()
 
+        if use_retry:
+            # Retry maximum 10 times, backoff on each retry
+            # Sleeps 1s, 2s, 4s, 8s, etc to a maximum of 120s between retries
+            # Retries on HTTP status codes 500, 502, 503, 504
+            retries = Retry(total=10,
+                            backoff_factor=1,
+                            status_forcelist=[500, 502, 503, 504])
+
+            # Use the retry strategy for all HTTPS requests
+            self.session.mount('https://', HTTPAdapter(max_retries=retries))
+
         # Keep resource classes as attributes of client.
         # Pass client to resource classes so resource object
         # can use the client.
@@ -71,7 +109,7 @@ class Client(object):
         self.Account = type('Account', (_Account,), attributes)
 
     def request(self, path, method='GET', params=None, data=None, files=None,
-                headers=None, raw=False, stream=False):
+                headers=None, raw=False, allow_redirects=True, stream=False):
         """
         Wrapper around requests.request()
 
@@ -91,27 +129,31 @@ class Client(object):
 
         headers['Accept'] = 'application/json'
 
+        if path.startswith('https://'):
+            url = path
+        else:
             url = BASE_URL + path
         log.debug('url: %s', url)
 
         response = self.session.request(
             method, url, params=params, data=data, files=files,
-            headers=headers, allow_redirects=True, stream=stream)
+            headers=headers, allow_redirects=allow_redirects, stream=stream)
         log.debug('response: %s', response)
         if raw:
             return response
 
         log.debug('content: %s', response.content)
         try:
-            response = json.loads(response.content)
+            body = json.loads(response.content.decode())
         except ValueError:
-            raise Exception('Server didn\'t send valid JSON:\n%s\n%s' % (
+            raise ServerError('InvalidJSON', response.content)
-                response, response.content))
 
-        if response['status'] == 'ERROR':
+        if body['status'] == 'ERROR':
-            raise Exception(response['error_type'])
+            log.error("API returned error: %s", body)
+            exception_class = {'4': ClientError, '5': ServerError}[str(response.status_code)[0]]
+            raise exception_class(body['error_type'], body['error_message'])
 
-        return response
+        return body
 
 
 class _BaseResource(object):
@@ -125,8 +167,8 @@ class _BaseResource(object):
         self.name = None
         self.__dict__.update(resource_dict)
         try:
-            self.created_at = parse(self.created_at)
+            self.created_at = strptime(self.created_at)
-        except AttributeError:
+        except Exception:
             self.created_at = None
 
     def __str__(self):
@@ -135,7 +177,7 @@ class _BaseResource(object):
     def __repr__(self):
         # shorten name for display
         name = self.name[:17] + '...' if len(self.name) > 20 else self.name
-        return '<%s id=%r, name="%r">' % (
+        return '<%s id=%r, name=%r>' % (
             self.__class__.__name__, self.id, name)
 
 
@@ -160,59 +202,113 @@ class _File(_BaseResource):
                 files = {'file': (name, f)}
             else:
                 files = {'file': f}
-            d = cls.client.request('/files/upload', method='POST',
+            d = cls.client.request(UPLOAD_URL, method='POST',
                                    data={'parent_id': parent_id}, files=files)
 
         f = d['file']
         return cls(f)
 
+    @classmethod
+    def upload_tus(cls, path, name=None, parent_id=0):
+        headers = {'Authorization': 'token %s' % cls.client.access_token}
+        metadata = {'parent_id': str(parent_id)}
+        if name:
+            metadata['name'] = name
+        with open(path) as f:
+            tus.upload(f, TUS_UPLOAD_URL, file_name=name, headers=headers, metadata=metadata)
+
     def dir(self):
         """List the files under directory."""
         return self.list(parent_id=self.id)
 
-    def download(self, dest='.', delete_after_download=False):
+    def download(self, dest='.', delete_after_download=False, chunk_size=CHUNK_SIZE):
         if self.content_type == 'application/x-directory':
-            self._download_directory(dest, delete_after_download)
+            self._download_directory(dest, delete_after_download, chunk_size)
         else:
-            self._download_file(dest, delete_after_download)
+            self._download_file(dest, delete_after_download, chunk_size)
 
-    def _download_directory(self, dest='.', delete_after_download=False):
+    def _download_directory(self, dest, delete_after_download, chunk_size):
-        name = self.name
+        name = _str(self.name)
-        if isinstance(name, unicode):
-            name = name.encode('utf-8', 'replace')
 
         dest = os.path.join(dest, name)
         if not os.path.exists(dest):
             os.mkdir(dest)
 
         for sub_file in self.dir():
-            sub_file.download(dest, delete_after_download)
+            sub_file.download(dest, delete_after_download, chunk_size)
 
         if delete_after_download:
             self.delete()
 
-    def _download_file(self, dest='.', delete_after_download=False):
+    def _verify_file(self, filepath):
-        response = self.client.request(
+        log.info('verifying crc32...')
-            '/files/%s/download' % self.id, raw=True, stream=True)
+        filesize = os.path.getsize(filepath)
+        if self.size != filesize:
+            logging.error('file %s is %d bytes, should be %s bytes' % (filepath, filesize, self.size))
+            return False
+
+        crcbin = 0
+        with open(filepath, 'rb') as f:
+            while True:
+                chunk = f.read(CHUNK_SIZE)
+                if not chunk:
+                    break
+
+                crcbin = binascii.crc32(chunk, crcbin) & 0xffffffff
 
-        filename = re.match(
+        crc32 = '%08x' % crcbin
-            'attachment; filename=(.*)',
-            response.headers['content-disposition']).groups()[0]
-        # If file name has spaces, it must have quotes around.
-        filename = filename.strip('"')
 
-        with open(os.path.join(dest, filename), 'wb') as f:
+        if crc32 != self.crc32:
-            for chunk in response.iter_content(chunk_size=1024):
+            logging.error('file %s CRC32 is %s, should be %s' % (filepath, crc32, self.crc32))
+            return False
+
+        return True
+
+    def _download_file(self, dest, delete_after_download, chunk_size):
+        name = _str(self.name)
+
+        filepath = os.path.join(dest, name)
+        if os.path.exists(filepath):
+            first_byte = os.path.getsize(filepath)
+
+            if first_byte == self.size:
+                log.warning('file %s exists and is the correct size %d' % (filepath, self.size))
+        else:
+            first_byte = 0
+
+        log.debug('file %s is currently %d, should be %d' % (filepath, first_byte, self.size))
+
+        if self.size == 0:
+            # Create an empty file
+            open(filepath, 'w').close()
+            log.debug('created empty file %s' % filepath)
+        else:
+            if first_byte < self.size:
+                with open(filepath, 'ab') as f:
+                    headers = {'Range': 'bytes=%d-' % first_byte}
+
+                    log.debug('request range: bytes=%d-' % first_byte)
+                    response = self.client.request('/files/%s/download' % self.id,
+                                                   headers=headers,
+                                                   raw=True,
+                                                   stream=True)
+
+                    for chunk in response.iter_content(chunk_size=chunk_size):
                         if chunk:  # filter out keep-alive new chunks
                             f.write(chunk)
-                    f.flush()
 
+        if self._verify_file(filepath):
             if delete_after_download:
                 self.delete()
 
     def delete(self):
         return self.client.request('/files/delete', method='POST',
-                                   data={'file_ids': str(self.id)})
+                                   data={'file_id': str(self.id)})
+
+    @classmethod
+    def delete_multi(cls, ids):
+        return cls.client.request('/files/delete', method='POST',
+                                  data={'file_ids': ','.join(map(str, ids))})
 
     def move(self, parent_id):
         return self.client.request('/files/move', method='POST',
@@ -239,6 +335,7 @@ class _Transfer(_BaseResource):
 
     @classmethod
     def add_url(cls, url, parent_id=0, extract=False, callback_url=None):
+        log.debug('callback_url is %s', callback_url) 
         d = cls.client.request('/transfers/add', method='POST', data=dict(
             url=url, save_parent_id=parent_id, extract=extract,
             callback_url=callback_url))
@@ -247,10 +344,10 @@ class _Transfer(_BaseResource):
 
     @classmethod
     def add_torrent(cls, path, parent_id=0, extract=False, callback_url=None):
-        with open(path) as f:
+        with open(path, 'rb') as f:
             files = {'file': f}
             d = cls.client.request('/files/upload', method='POST', files=files,
-                                   data=dict(save_parent_id=parent_id,
+                                   data=dict(parent_id=parent_id,
                                              extract=extract,
                                              callback_url=callback_url))
         t = d['transfer']
@@ -260,6 +357,17 @@ class _Transfer(_BaseResource):
     def clean(cls):
         return cls.client.request('/transfers/clean', method='POST')
 
+    def cancel(self):
+        return self.client.request('/transfers/cancel',
+                                   method='POST',
+                                   data={'transfer_ids': self.id})
+
+    @classmethod
+    def cancel_multi(cls, ids):
+        return cls.client.request('/transfers/cancel',
+                                  method='POST',
+                                  data={'transfer_ids': ','.join(map(str, ids))})
+
 
 class _Account(_BaseResource):
 
@@ -270,3 +378,31 @@ class _Account(_BaseResource):
     @classmethod
     def settings(cls):
         return cls.client.request('/account/settings', method='GET')
+
+
+# Due to a nasty bug in datetime module, datetime.strptime calls
+# are not thread-safe and can throw a TypeError. Details: https://bugs.python.org/issue7980
+# Here we are implementing simple RFC3339 parser which is used in Put.io APIv2.
+def strptime(date):
+    """Returns datetime object from the given date, which is in a specific format: YYYY-MM-ddTHH:mm:ss"""
+    d = {
+        'year': date[0:4],
+        'month': date[5:7],
+        'day': date[8:10],
+        'hour': date[11:13],
+        'minute': date[14:16],
+        'second': date[17:],
+    }
+
+    d = dict((k, int(v)) for k, v in d.iteritems())
+    return datetime(**d)
+
+
+def _str(s):
+    """Python 3 compatibility function for converting to str."""
+    try:
+        if isinstance(s, unicode):
+            return s.encode('utf-8', 'replace')
+    except NameError:
+        pass
+    return s