linaro-cp.py: Convert to python-requests.

Change-Id: I9a61d388e7182773791fda0ffeb953352e8cf0f1

1 files changed, 44 insertions, 92 deletions

diff --git a/linaro-cp.py b/linaro-cp.py
index 066cfe5..ac60b5b 100755
--- a/linaro-cp.py
+++ b/linaro-cp.py
@@ -2,16 +2,17 @@
 
 import argparse
 import atexit
-import cStringIO
 import os
 import mimetypes
-import pycurl
 import sys
 import tempfile
 import time
 import re
 import pprint
 
+import requests
+
+
 # Public artifacts BUILD-INFO.txt
 build_info = 'Format-Version: 0.5\n\nFiles-Pattern: *\nLicense-Type: open\n'
 
@@ -22,44 +23,26 @@ class API_v1(object):
         self.api_base = server
         self.build_info = build_info
         self.api_key = api_key
-        self.curl = pycurl.Curl()
-
-    def __del__(self):
-        self.curl.close()
-
-    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:
-            if retry_count > 0:
-                # server could be reloading or something. give it a second and
-                # try again
-                print('Upload failed for %s, retrying in 2 seconds' % url)
-                time.sleep(2)
-                return self._upload_data(url, data, headers, retry_count - 1)
-            else:
-                return repr(e)
-
-        code = self.curl.getinfo(pycurl.RESPONSE_CODE)
-        if code == 503 and retry_count > 0:
-                print('503 failure for %s, retrying in 2 seconds' % url)
-                time.sleep(2)
-                return self._upload_data(url, data, headers, retry_count - 1)
-        if code not in (200, 201):
-            return 'HTTP_%d: %s' % (code, response.getvalue())
+
+    def _upload_data(self, url, data=None, files=None, headers=None, retry_count=3):
+        while retry_count:
+            try:
+                resp = requests.post(url, headers=headers, data=data, files=files)
+                if resp.status_code in (200, 201):
+                    return resp
+                else:
+                    print("Unsuccessful status:", resp.status_code)
+            except Exception as e:
+                print(e)
+            print('Upload failed for %s, retrying in 2 seconds' % url)
+            time.sleep(2)
+            retry_count -= 1
+
+        raise Exception("LLP request failed")
 
     def upload_file(self, url, filename):
-        data = [
-            ('file', (pycurl.FORM_FILE, filename)),
-            ('key', (pycurl.FORM_CONTENTS, self.api_key)),
-        ]
-        return self._upload_data(url, data)
+        with open(filename, 'rb') as f:
+            return self._upload_data(url, files={'file': f}, data={'key': self.api_key})
 
     def upload_transfer_queue(self, transfer_queue):
         transfer_failures = []
@@ -135,16 +118,12 @@ class API_v2(API_v1):
         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)
+        with open(filename, 'rb') as f:
+            return self._upload_data(url, files={'file': f}, headers={'AuthToken': self.api_key})
 
     def link_latest(self, dst):
-        headers = ['AuthToken: ' + self.api_key]
         url = self.server + '/api/v2/link_latest/' + dst
-        # pycurl requires data to be passed, or it will do an
-        # HTTP GET even though we said to POST
-        return self._upload_data(url, [('foo', 'bar')], headers)
+        self._upload_data(url, headers={'AuthToken': self.api_key})
 
 
 class API_v3(API_v1):
@@ -152,42 +131,24 @@ class API_v3(API_v1):
         super(API_v3, self).__init__(server, build_info, api_key)
         self.api_base = server + '/api/v3/publish/'
 
-    def _upload_data(self, url, data, headers=None, retry_count=3):
-        self.last_headers = cStringIO.StringIO()
-        self.curl.setopt(pycurl.HEADERFUNCTION, self.last_headers.write)
-        return super(API_v3, self)._upload_data(
-            url, data, headers, retry_count)
-
     def _put_s3(self, url, filename, mtype, retry_count=3):
-        response = cStringIO.StringIO()
         size = os.path.getsize(filename)
-        headers = ['Content-Type: ' + mtype]
-        c = pycurl.Curl()
-        c.setopt(pycurl.URL, url)
-        c.setopt(pycurl.HTTPHEADER, headers)
-        c.setopt(pycurl.INFILESIZE, size)
-        c.setopt(pycurl.PUT, 1)
-        c.setopt(pycurl.WRITEFUNCTION, response.write)
-        try:
-            with open(filename, 'rb') as f:
-                c.setopt(pycurl.INFILE, f)
-                c.perform()
-        except Exception as e:
-            if retry_count > 0:
-                # server could be reloading or something. give it a second and
-                # try again
-                print('Upload failed for %s, retrying in 2 seconds' % url)
-                time.sleep(2)
-                return self._put_s3(url, filename, mtype, retry_count - 1)
-            else:
-                return repr(e)
-        code = c.getinfo(pycurl.RESPONSE_CODE)
-        if code not in (200, 201):
-            if retry_count > 0:
-                print('%d failure for %s, retrying in 2 seconds' % (code, url))
-                time.sleep(2)
-                return self._put_s3(url, filename, mtype, retry_count - 1)
-            return response.getvalue()
+        headers = {'Content-Type': mtype, 'Content-Length': str(size)}
+        while retry_count:
+            try:
+                with open(filename, 'rb') as f:
+                    resp = requests.put(url, headers=headers, data=f)
+                if resp.status_code in (200, 201):
+                    return
+                else:
+                    print("Unsuccessful status:", resp.status_code)
+            except Exception as e:
+                print(e)
+            print('Upload failed for %s, retrying in 2 seconds' % url)
+            time.sleep(2)
+            retry_count -= 1
+
+        return "S3 upload failed"
 
     def upload_file(self, url, filename):
         # ask llp for an s3 tempurl:
@@ -195,24 +156,15 @@ class API_v3(API_v1):
         if not mtype:
             mtype = 'other'
 
-        headers = ['AuthToken: ' + self.api_key]
-        code = self._upload_data(url, [('Content-Type', mtype)], headers)
-        if code:
-            return code
+        resp = self._upload_data(url, data={'Content-Type': mtype}, headers={'AuthToken': self.api_key})
 
         # now find the tempurl and really publish
-        for header in self.last_headers.getvalue().split('\n'):
-            if header.startswith('Location:'):
-                location = header[9:].strip()
-                return self._put_s3(location, filename, mtype)
-        raise RuntimeError('l-l-p response missing s3 location')
+        location = resp.headers['location']
+        return self._put_s3(location, filename, mtype)
 
     def link_latest(self, dst):
-        headers = ['AuthToken: ' + self.api_key]
         url = self.server + '/api/v3/link_latest/' + dst
-        # pycurl requires data to be passed, or it will do an
-        # HTTP GET even though we said to POST
-        return self._upload_data(url, [('foo', 'bar')], headers)
+        self._upload_data(url, headers={'AuthToken': self.api_key})
 
 
 def main():