bug 738000: enable CSRF protection in existing balrog code. r=rail

mozilla · Jun 5, 2012 · c4e0697 · c4e0697
1 parent 3524f2b
commit c4e0697
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 13 deletions.
diff --git a/lib/python/balrog/client/api.py b/lib/python/balrog/client/api.py
@@ -1,18 +1,19 @@
 # TODO: extend API to handle release blobs
-
 import logging
 import requests
 import os
 
 CA_BUNDLE = os.path.join(os.path.dirname(__file__),
                         '../../../../misc/certs/GeoTrust_Global_CA.crt')
 
+def is_csrf_token_expired(token):
+    from datetime import datetime
+    expiry = token.split('##')[0]
+    if expiry <= datetime.now().strftime('%Y%m%d%H%M%S'):
+        return True
+    return False
 
 class API(object):
-
-    url_template = \
-        '%(api_root)s/releases/%(name)s/builds/%(build_target)s/%(locale)s'
-
     verify = False
     auth = None
 
@@ -39,22 +40,59 @@ def __init__(self, api_root='https://balrog.build.mozilla.org',
         self.auth = auth
         self.timeout = timeout
         self.config = dict(danger_mode=raise_exceptions)
+        self.session = requests.session()
+        self.csrf_token = None
+
+    def request(self, data=None, method='GET', url_template_vars={}):
+        url = self.api_root + self.url_template % url_template_vars
+        # If we'll be modifying things, do a GET first to get a CSRF token
+        # and possibly a data_version.
+        if method != 'GET' and method != 'HEAD':
+            # Use the URL of the resource we're going to modify first,
+            # because we'll need its data_version if it exists.
+            try:
+                res = self.do_request(url, None, 'HEAD', {})
+                data['data_version'] = res.headers['X-Data-Version']
+                # We may already have a non-expired CSRF token, but it's
+                # faster/easier just to set it again even if we do, since
+                # we've already made the request.
+                data['csrf_token'] = self.csrf_token = res.headers['X-CSRF-Token']
+            except requests.HTTPError, e:
+                # However, if the resource doesn't exist yet we may as well
+                # not bother doing another request solely for a token unless
+                # we don't have a valid one already.
+                if e.response.status_code != 404:
+                    raise
+                if not self.csrf_token or is_csrf_token_expired(self.csrf_token):
+                    res = self.do_request(self.api_root + '/csrf_token', None, 'HEAD', {})
+                    data['csrf_token'] = self.csrf_token = res.headers['X-CSRF-Token']
 
-    def request(self, url, data=None, method='GET'):
+            logging.debug('Got CSRF Token: %s' % self.csrf_token)
+        return self.do_request(url, data, method, url_template_vars)
+
+    def do_request(self, url, data, method, url_template_vars):
         logging.debug('Balrog request to %s' % url)
         logging.debug('Data sent: %s' % data)
-        return requests.request(method=method, url=url, data=data,
+        try:
-                                config=self.config, timeout=self.timeout,
+            return self.session.request(method=method, url=url, data=data,
-                                verify=self.verify, auth=self.auth)
+                                        config=self.config, timeout=self.timeout,
+                                        verify=self.verify, auth=self.auth)
+        except requests.HTTPError, e:
+            logging.error('Caught HTTPError: %s' % e.response.content)
+            raise
+
+
+class SingleLocale(API):
+    url_template = '/releases/%(name)s/builds/%(build_target)s/%(locale)s'
 
     def update_build(self, name, product, version, build_target, locale,
                      buildData, copyTo=None):
         url_template_vars = dict(api_root=self.api_root, name=name,
                                  locale=locale, build_target=build_target)
-        url = self.url_template % url_template_vars
         data = dict(product=product, version=version,
                     data=buildData)
         if copyTo:
             data['copyTo'] = copyTo
 
-        return self.request(method='PUT', url=url, data=data)
+        return self.request(method='PUT', data=data,
+                            url_template_vars=url_template_vars)
diff --git a/lib/python/balrog/client/cli.py b/lib/python/balrog/client/cli.py
@@ -4,7 +4,7 @@
     import json
 
 from release.platforms import buildbot2updatePlatforms
-from balrog.client.api import API
+from balrog.client.api import SingleLocale
 
 
 def get_nightly_blob_name(appName, branch, build_type, suffix, dummy=False):
@@ -70,7 +70,7 @@ def generate_data(self):
     def run(self):
         data = self.generate_data()
         data = json.dumps(data)
-        api = API(auth=self.auth, api_root=self.api_root)
+        api = SingleLocale(auth=self.auth, api_root=self.api_root)
         copyTo = [get_nightly_blob_name(
             self.appName, self.branch, self.build_type, 'latest', self.dummy)]
         copyTo = json.dumps(copyTo)