Merge pull request #395 from qiniu/enhance

Enhance

Author: mandytong002

qiniu/auth.py

@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
-
+import base64
 import hmac
 import time
 from hashlib import sha1
@@ -158,6 +158,16 @@ def upload_token(
 
         return self.__upload_token(args)
 
+    @staticmethod
+    def up_token_decode(up_token):
+        up_token_list = up_token.split(':')
+        ak = up_token_list[0]
+        sign = base64.urlsafe_b64decode(up_token_list[1])
+        decode_policy = base64.urlsafe_b64decode(up_token_list[2])
+        decode_policy = decode_policy.decode('utf-8')
+        dict_policy = json.loads(decode_policy)
+        return ak, sign, dict_policy
+
     def __upload_token(self, policy):
         data = json.dumps(policy, separators=(',', ':'))
         return self.token_with_data(data)

qiniu/services/storage/uploader.py

@@ -1,11 +1,10 @@
 # -*- coding: utf-8 -*-
-
 import hashlib
 import json
 import os
 import time
 
-from qiniu import config
+from qiniu import config, Auth
 from qiniu.utils import urlsafe_base64_encode, crc32, file_crc32, _file_iter, rfc_from_timestamp
 from qiniu import http
 from .upload_progress_recorder import UploadProgressRecorder
@@ -239,6 +238,7 @@ def upload(self):
         self.recovery_index = 1
         self.expiredAt = None
         self.uploadId = None
+        self.get_bucket()
         host = self.get_up_host()
         if self.version == 'v1':
             offset = self.recovery_from_record()
@@ -388,3 +388,9 @@ def put(self, url, data, headers):
 
     def get_parts(self, block_status):
         return sorted(block_status, key=lambda i: i['partNumber'])
+
+    def get_bucket(self):
+        if self.bucket_name is None or self.bucket_name == '':
+            _, _, policy = Auth.up_token_decode(self.up_token)
+            if policy != {}:
+                self.bucket_name = policy['scope'].split(':')[0]

qiniu/utils.py

@@ -90,6 +90,7 @@ def _file_iter(input_stream, size, offset=0):
     while d:
         yield d
         d = input_stream.read(size)
+    input_stream.seek(0)
 
 
 def _sha1(data):

test_qiniu.py

@@ -390,6 +390,17 @@ def test_put_stream(self):
                                    self.mime_type, part_size=None, version=None, bucket_name=None)
             assert ret['key'] == key
 
+    def test_put_stream_v2_without_bucket_name(self):
+        localfile = __file__
+        key = 'test_file_r'
+        size = os.stat(localfile).st_size
+        set_default(default_zone=Zone('http://upload.qiniup.com'))
+        with open(localfile, 'rb') as input_stream:
+            token = self.q.upload_token(bucket_name, key)
+            ret, info = put_stream(token, key, input_stream, os.path.basename(__file__), size, hostscache_dir,
+                                   self.params,
+                                   self.mime_type, part_size=1024 * 1024 * 10, version='v2')
+            assert ret['key'] == key
 
     def test_put_2m_stream_v2(self):
         localfile = create_temp_file(2 * 1024 * 1024 + 1)