104 lines
3.2 KiB
Python
104 lines
3.2 KiB
Python
# Copyright 2023 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
'Utility functions for API requests and responses.'
|
|
|
|
import itertools
|
|
import json
|
|
import logging
|
|
import re
|
|
|
|
from . import HTTPRequest, PluginError
|
|
|
|
MP_PART = '''\
|
|
Content-Type: application/http
|
|
MIME-Version: 1.0
|
|
Content-Transfer-Encoding: binary
|
|
|
|
GET {}?alt=json HTTP/1.1
|
|
Content-Type: application/json
|
|
MIME-Version: 1.0
|
|
Content-Length: 0
|
|
Accept: application/json
|
|
Accept-Encoding: gzip, deflate
|
|
Host: compute.googleapis.com
|
|
|
|
'''
|
|
RE_URL = re.compile(r'pageToken=[^&]+&?')
|
|
|
|
|
|
def batched(iterable, n):
|
|
'Batches data into lists of length n. The last batch may be shorter.'
|
|
# batched('ABCDEFG', 3) --> ABC DEF G
|
|
if n < 1:
|
|
raise ValueError('n must be at least one')
|
|
it = iter(iterable)
|
|
while (batch := list(itertools.islice(it, n))):
|
|
yield batch
|
|
|
|
|
|
def parse_cai_results(data, name, resource_type=None, method='search'):
|
|
'Parses an asset API response and returns individual results.'
|
|
results = data.get('results' if method == 'search' else 'assets')
|
|
if not results:
|
|
logging.info(f'no results for {name}')
|
|
return
|
|
for result in results:
|
|
if resource_type and result['assetType'] != resource_type:
|
|
logging.warn(f'result for wrong type {result["assetType"]}')
|
|
continue
|
|
yield result
|
|
|
|
|
|
def parse_page_token(data, url):
|
|
'Detect next page token in result and return next page URL.'
|
|
page_token = data.get('nextPageToken')
|
|
if page_token:
|
|
logging.info(f'page token {page_token}')
|
|
if 'pageToken' in url:
|
|
return RE_URL.sub(f'pageToken={page_token}', url)
|
|
else:
|
|
return f'{url}&pageToken={page_token}'
|
|
|
|
|
|
def poor_man_mp_request(urls, boundary='1234567890'):
|
|
'Bundles URLs into a single multipart mixed batched request.'
|
|
boundary = f'--{boundary}'
|
|
data = [boundary]
|
|
for url in urls:
|
|
data += ['\n', MP_PART.format(url), boundary]
|
|
data.append('--\n')
|
|
headers = {'content-type': f'multipart/mixed; boundary={boundary[2:]}'}
|
|
return HTTPRequest('https://compute.googleapis.com/batch/compute/v1', headers,
|
|
''.join(data), False)
|
|
|
|
|
|
def poor_man_mp_response(content_type, content):
|
|
'Parses a multipart mixed response and returns individual parts.'
|
|
try:
|
|
_, boundary = content_type.split('=')
|
|
except ValueError:
|
|
raise PluginError('no boundary found in content type')
|
|
content = content.decode('utf-8').strip()[:-2]
|
|
if boundary not in content:
|
|
raise PluginError('MIME boundary not found')
|
|
for part in content.split(f'--{boundary}'):
|
|
part = part.strip()
|
|
if not part:
|
|
continue
|
|
try:
|
|
mime_header, header, body = part.split('\r\n\r\n', 3)
|
|
except ValueError:
|
|
raise PluginError('cannot parse MIME part')
|
|
yield json.loads(body)
|