amancevice · amancevice · May 6, 2020 · Feb 20, 2020
diff --git a/Makefile b/Makefile
@@ -0,0 +1,19 @@
+.PHONY: default clean test up validate
+
+default: test validate
+
+.terraform:
+	terraform init
+
+clean:
+	rm -rf .terraform
+
+test:
+	flake8
+	pytest
+
+up:
+	lambda-gateway index.proxy_request -B simple
+
+validate: | .terraform
+	terraform validate
diff --git a/index.py b/index.py
@@ -2,6 +2,8 @@
 import os
 import string
 import sys
+from distutils.version import StrictVersion
+from xml.etree import ElementTree as xml
 
 import boto3
 
@@ -11,9 +13,31 @@
     '<!DOCTYPE html><html><head><title>$title</title></head>'
     '<body><h1>$title</h1>$anchors</body></html>'
 )
+SEARCH = string.Template(
+    "<?xml version='1.0'?>"
+    '<methodResponse><params><param><value><array><data>'
+    '$data'
+    '</data></array></value></param></params></methodResponse>'
+)
+SEARCH_VALUE = string.Template(
+    '<struct>'
+    '<member><name>name</name><value><string>'
+    '$name'
+    '</string></value></member>'
+    '<member><name>summary</name><value><string>'
+    '$summary'
+    '</string></value></member>'
+    '<member><name>version</name><value><string>'
+    '$version'
+    '</string></value></member>'
+    '<member><name>_pypi_ordering</name><value><boolean>'
+    '0'
+    '</boolean></value></member>'
+    '</struct>'
+)
 
 S3 = boto3.client('s3')
-S3_BUCKET = os.getenv('S3_BUCKET')
+S3_BUCKET = os.getenv('S3_BUCKET', 'serverless-pypi')
 S3_PAGINATOR = S3.get_paginator('list_objects')
 S3_PRESIGNED_URL_TTL = int(os.getenv('S3_PRESIGNED_URL_TTL', '900'))
 
@@ -79,7 +103,7 @@ def get_package_index(name):
     return resp
 
 
-def get_response(path):
+def get_response(path, *_):
     """ GET /{BASE_PATH}/*
 
         :param str path: Request path
@@ -101,7 +125,7 @@ def get_response(path):
     return reject(403, message='Forbidden')
 
 
-def head_response(path):
+def head_response(path, *_):
     """ HEAD /{BASE_PATH}/*
 
         :param str path: Request path
@@ -113,6 +137,19 @@ def head_response(path):
     return res
 
 
+def post_response(path, body):
+    """ POST /{BASE_PATH}/
+
+        :param str path: POST path
+        :param str body: POST body
+        :return dict: Response
+    """
+    if path == BASE_PATH:
+        return search(body)
+
+    return reject(403, message='Forbidden')
+
+
 def presign(key):
     """ Presign package URLs.
 
@@ -181,11 +218,41 @@ def reject(status_code, **kwargs):
     return res
 
 
+def search(request):
+    """ Search for pips.
+
+        :param str request: XML request string
+        :return str: XML response
+    """
+    tree = xml.fromstring(request)
+    term = tree.find('.//string').text  # TODO this is not ideal
+
+    hits = {}
+    for page in S3_PAGINATOR.paginate(Bucket=S3_BUCKET):
+        for obj in page.get('Contents'):
+            key = obj.get('Key')
+            if term in key and 'index.html' != key:
+                *_, name, tarball = key.split('/')
+                _, version = tarball.replace('.tar.gz', '').split(f'{name}-')
+                version = StrictVersion(version)
+                if name not in hits or hits[name]['version'] < version:
+                    hits[name] = {
+                        'name': name,
+                        'version': version,
+                        'summary': f's3://{S3_BUCKET}/{key}',
+                    }
+    data = [SEARCH_VALUE.safe_substitute(**x) for x in hits.values()]
+    body = SEARCH.safe_substitute(data=''.join(data))
+    resp = proxy_reponse(body, 'text/xml')
+    return resp
+
+
 # Lambda handlers
 
 ROUTES = {
     'GET': get_response,
     'HEAD': head_response,
+    'POST': post_response,
 }
 
 
@@ -197,10 +264,11 @@ def proxy_request(event, *_):
     # Get HTTP request method/path
     method = event.get('httpMethod')
     path = event.get('path').strip('/')
+    body = event.get('body')
 
     # Get HTTP response
     try:
-        res = ROUTES[method](path)
+        res = ROUTES[method](path, body)
     except KeyError:
         res = reject(403, message='Forbidden')
 

diff --git a/index_test.py b/index_test.py
@@ -1,5 +1,6 @@
 import io
 import json
+import textwrap
 from unittest import mock
 
 import pytest
@@ -149,6 +150,17 @@ def test_proxy_request_get(mock_idx):
     mock_idx.assert_called_once_with()
 
 
+@mock.patch('index.search')
+def test_proxy_reponse_post(mock_search):
+    mock_search.return_value = index.proxy_reponse('{}')
+    index.proxy_request({
+        'body': '<SEARCH_XML>',
+        'httpMethod': 'POST',
+        'path': '/simple/',
+    })
+    mock_search.assert_called_once_with('<SEARCH_XML>')
+
+
 @mock.patch('index.get_package_index')
 def test_proxy_request_get_package(mock_pkg):
     mock_pkg.return_value = index.proxy_reponse(PACKAGE_INDEX)
@@ -159,7 +171,8 @@ def test_proxy_request_get_package(mock_pkg):
 @pytest.mark.parametrize('http_method,path,status_code', [
     ('HEAD', '/fizz/buzz/jazz', 403),
     ('GET', '/fizz/buzz/jazz', 403),
-    ('POST', '/fizz/buzz', 403),
+    ('POST', '/fizz/buzz/jazz', 403),
+    ('OPTIONS', '/fizz/buzz/jazz', 403),
 ])
 def test_proxy_request_reject(http_method, path, status_code):
     ret = index.proxy_request({'httpMethod': http_method, 'path': path})
@@ -178,3 +191,56 @@ def test_reindex_bucket():
         Key='index.html',
         Body=SIMPLE_INDEX.encode(),
     )
+
+
+@pytest.mark.parametrize('pip', ['fizz'])
+def test_search(pip):
+    index.S3_PAGINATOR.paginate.return_value = iter(S3_INDEX_RESPONSE)
+    request = textwrap.dedent(f'''\
+        <?xml version='1.0'?>
+        <methodCall>
+        <methodName>search</methodName>
+        <params>
+        <param>
+        <value><struct>
+        <member>
+        <name>name</name>
+        <value><array><data>
+        <value><string>{pip}</string></value>
+        </data></array></value>
+        </member>
+        <member>
+        <name>summary</name>
+        <value><array><data>
+        <value><string>{pip}</string></value>
+        </data></array></value>
+        </member>
+        </struct></value>
+        </param>
+        <param>
+        <value><string>or</string></value>
+        </param>
+        </params>
+        </methodCall>\
+    ''')
+    body = (
+        "<?xml version='1.0'?><methodResponse><params><param><value><array>"
+        "<data><struct><member><name>name</name><value><string>fizz</string>"
+        "</value></member><member><name>summary</name><value>"
+        "<string>s3://serverless-pypi/simple/fizz/fizz-1.2.3.tar.gz</string>"
+        "</value></member><member><name>version</name><value>"
+        "<string>1.2.3</string></value></member><member>"
+        "<name>_pypi_ordering</name><value><boolean>0</boolean></value>"
+        "</member></struct></data></array></value></param></params>"
+        "</methodResponse>"
+    )
+    ret = index.search(request)
+    exp = {
+        'body': body,
+        'statusCode': 200,
+        'headers': {
+            'Content-Size': len(body),
+            'Content-Type': 'text/xml; charset=UTF-8',
+        },
+    }
+    assert ret == exp