Add whitelist check for github head version packages

compatibility_lib/compatibility_lib/compatibility_checker.py

@@ -22,7 +22,13 @@
 
 from compatibility_lib import configs
 
-SERVER_URL = 'http://104.197.8.72'
+SERVER_URL = 'http://0.0.0.0:8888'
+
+PACKAGE_NOT_IN_WHITELIST = 'Request contains third party github head packages.'
+
+UNKNOWN_STATUS_RESULT = {
+    'result': 'UNKNOWN',
+}
 
 
 class CompatibilityChecker(object):
@@ -37,8 +43,13 @@ def check(self, packages, python_version):
             'package': packages
         }
         result = requests.get(SERVER_URL, params=data)
+        content = result.content.decode('utf-8')
+        if content == PACKAGE_NOT_IN_WHITELIST:
+            UNKNOWN_STATUS_RESULT['packages'] = packages
+            UNKNOWN_STATUS_RESULT['description'] = PACKAGE_NOT_IN_WHITELIST
+            return UNKNOWN_STATUS_RESULT
 
-        return json.loads(result.content.decode('utf-8'))
+        return json.loads(content)
 
     @retrying.retry(wait_exponential_multiplier=5000,
                     wait_exponential_max=20000)

compatibility_server/compatibility_checker_server.py

@@ -49,6 +49,12 @@
 
 import pip_checker
 
+# White list Google owned Python packages
+GITHUB_PREFIX = 'github.com/'
+WHITELIST_GITHUB_REPO = ['GoogleCloudPlatform/',
+                         'google/',
+                         'googleapis/']
+
 
 def _parse_python_version_to_interpreter_mapping(s):
     version_to_interpreter = {}
@@ -96,6 +102,13 @@ def _check(self, start_response, python_version, packages):
                            [('Content-Type', 'text/plain; charset=utf-8')])
             return [b'Request must specify at least one package']
 
+        sanitized_packages = self._sanitize_packages(packages)
+
+        if sanitized_packages != packages:
+            start_response('400 Bad Request',
+                           [('Content-Type', 'text/plain; charset=utf-8')])
+            return [b'Request contains third party github head packages.']
+
         if not python_version:
             start_response('400 Bad Request',
                            [('Content-Type', 'text/plain; charset=utf-8')])
@@ -139,6 +152,21 @@ def _check(self, start_response, python_version, packages):
         start_response('200 OK', [('Content-Type', 'application/json')])
         return [json.dumps(results).encode('utf-8')]
 
+    def _sanitize_packages(self, packages):
+        # If checking github head version, only run checks for whitelisted
+        # repos.
+        sanitized_packages = []
+        for pkg in packages:
+            if GITHUB_PREFIX in pkg:
+                for whitelist_repo in WHITELIST_GITHUB_REPO:
+                    github_whitelist = GITHUB_PREFIX + whitelist_repo
+                    if github_whitelist in pkg:
+                        sanitized_packages.append(pkg)
+            else:
+                sanitized_packages.append(pkg)
+
+        return sanitized_packages
+
     def _wsgi_app(self, environ, start_response):
         if environ.get('REQUEST_METHOD') == 'GET':
             parameters = urllib.parse.parse_qs(environ.get('QUERY_STRING', ''))
@@ -170,8 +198,7 @@ def _wsgi_app(self, environ, start_response):
                 environ.get('REQUEST_METHOD').encode('utf-8')
             ]
 
-        return self._check(start_response, python_version,
-                           packages)
+        return self._check(start_response, python_version, packages)
 
     def serve(self):
         with wsgiref.simple_server.make_server(self._host, self._port,