add support for v2 api

This allows you to pass a new APIv2 token and use the updated
publishing API.

Change-Id: I53ecc742af5c036d18e6e4b0c0de29e63db69292

1 files changed, 32 insertions, 7 deletions

diff --git a/linaro-cp.py b/linaro-cp.py
index e7d799f..2ed948f 100755
--- a/linaro-cp.py
+++ b/linaro-cp.py
@@ -24,11 +24,13 @@ class API_v1(object):
     def __del__(self):
         self.curl.close()
 
-    def _upload_data(self, url, data, retry_count=3):
+    def _upload_data(self, url, data, headers=None, retry_count=3):
         response = cStringIO.StringIO()
         self.curl.setopt(pycurl.URL, url)
         self.curl.setopt(pycurl.HTTPPOST, data)
         self.curl.setopt(pycurl.WRITEFUNCTION, response.write)
+        if headers:
+            self.curl.setopt(pycurl.HTTPHEADER, headers)
         try:
             self.curl.perform()
         except Exception as e:
@@ -41,7 +43,9 @@ class API_v1(object):
             else:
                 return str(e)
 
-        return response.getvalue()
+        code = self.curl.getinfo(pycurl.RESPONSE_CODE)
+        if code not in (200, 201):
+            return response.getvalue()
 
     def upload_file(self, url, filename):
         data = [
@@ -53,10 +57,10 @@ class API_v1(object):
     def upload_transfer_queue(self, transfer_queue):
         transfer_failures = []
         for transfer_item in transfer_queue:
-            http_status = self.upload_file(
+            failure = self.upload_file(
                 transfer_item, transfer_queue[transfer_item])
-            if http_status != 'OK':
-                transfer_failures.append('%s: %s' % (transfer_item, http_status))
+            if failure:
+                transfer_failures.append('%s: %s' % (transfer_item, failure))
         return transfer_failures
 
     def get_transfer_queue(self, src, dst):
@@ -85,10 +89,23 @@ class API_v1(object):
         return self.upload_transfer_queue(transfer_queue)
 
 
+class API_v2(API_v1):
+    def __init__(self, server, build_info, api_key):
+        super(API_v2, self).__init__(server, build_info, api_key)
+        self.api_base = server + 'api/v2/publish/'
+
+    def upload_file(self, url, filename):
+        headers = ['AuthToken: ' + self.api_key]
+        data = [('file', (pycurl.FORM_FILE, filename))]
+        return self._upload_data(url, data, headers)
+
+
 def main():
     parser = argparse.ArgumentParser(
         description='Copy file(s) from source to destination')
     parser.add_argument('-k', '--key', help='key used for the copy')
+    parser.add_argument('-a', '--api_version', choices=('1', '2'), default='1',
+                        help='API version to use. default=%(default)s')
     parser.add_argument('--server', default='http://snapshots.linaro.org/',
                         help='Publishing API server. default=%(default)s')
     parser.add_argument('-b', '--build-info', help='Custom build-info file')
@@ -103,14 +120,22 @@ def main():
     else:
         key = os.environ.get('PUBLISH_KEY')
         if key is None:
-            sys.exit('Key is not defined.')
+            key = os.environ.get('PUBLISH_TOKEN')
+            if key:
+                arguments.api_version = '2'
+            else:
+                sys.exit('Key is not defined.')
 
     if not arguments.build_info:
         fd, arguments.build_info = tempfile.mkstemp(prefix='BUILD-INFO.txt')
         atexit.register(os.unlink, arguments.build_info)
         os.write(fd, build_info)
 
-    api = API_v1(arguments.server, arguments.build_info, key)
+    if arguments.api_version == '1':
+        api = API_v1(arguments.server, arguments.build_info, key)
+    else:
+        api = API_v2(arguments.server, arguments.build_info, key)
+
     transfer_failures = api.upload(arguments.src, arguments.dst)
 
     if len(transfer_failures) > 0: