Extract image download method into a mixin

This is the same code for v1 and v2.

While we're in there, add checksum verification for when download
is used with an output file.

Change-Id: I35675fdbc29728b39ca76fc411f656e6234623a5
This commit is contained in:
Artem Goncharov 2019-05-11 09:40:38 +02:00
parent 1e810595c6
commit cf9922e885
6 changed files with 117 additions and 131 deletions

View File

@ -0,0 +1,85 @@
# 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.
import io
import hashlib
import six
from openstack import exceptions
from openstack import utils
def _verify_checksum(md5, checksum):
if checksum:
digest = md5.hexdigest()
if digest != checksum:
raise exceptions.InvalidResponse(
"checksum mismatch: %s != %s" % (checksum, digest))
class DownloadMixin(object):
def download(self, session, stream=False, output=None, chunk_size=1024):
"""Download the data contained in an image"""
# TODO(briancurtin): This method should probably offload the get
# operation into another thread or something of that nature.
url = utils.urljoin(self.base_path, self.id, 'file')
resp = session.get(url, stream=stream)
# See the following bug report for details on why the checksum
# code may sometimes depend on a second GET call.
# https://storyboard.openstack.org/#!/story/1619675
checksum = resp.headers.get("Content-MD5")
if checksum is None:
# If we don't receive the Content-MD5 header with the download,
# make an additional call to get the image details and look at
# the checksum attribute.
details = self.fetch(session)
checksum = details.checksum
md5 = hashlib.md5()
if output:
try:
# In python 2 we might get StringIO - delete it as soon as
# py2 support is dropped
if isinstance(output, io.IOBase) \
or isinstance(output, six.StringIO):
for chunk in resp.iter_content(chunk_size=chunk_size):
output.write(chunk)
md5.update(chunk)
else:
with open(output, 'wb') as fd:
for chunk in resp.iter_content(
chunk_size=chunk_size):
fd.write(chunk)
md5.update(chunk)
_verify_checksum(md5, checksum)
return resp
except Exception as e:
raise exceptions.SDKException(
"Unable to download image: %s" % e)
# if we are returning the repsonse object, ensure that it
# has the content-md5 header so that the caller doesn't
# need to jump through the same hoops through which we
# just jumped.
if stream:
resp.headers['content-md5'] = checksum
return resp
if checksum is not None:
_verify_checksum(hashlib.md5(resp.content), checksum)
else:
session.log.warn(
"Unable to verify the integrity of image %s", (self.id))
return resp

View File

@ -9,16 +9,11 @@
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
import hashlib
import io
import six
from openstack import exceptions
from openstack.image import _download
from openstack import resource
from openstack import utils
class Image(resource.Resource):
class Image(resource.Resource, _download.DownloadMixin):
resource_key = 'image'
resources_key = 'images'
base_path = '/images'
@ -78,58 +73,3 @@ class Image(resource.Resource):
status = resource.Body('status')
#: The timestamp when this image was last updated.
updated_at = resource.Body('updated_at')
def download(self, session, stream=False, output=None, chunk_size=1024):
"""Download the data contained in an image"""
# TODO(briancurtin): This method should probably offload the get
# operation into another thread or something of that nature.
url = utils.urljoin(self.base_path, self.id, 'file')
resp = session.get(url, stream=stream)
# See the following bug report for details on why the checksum
# code may sometimes depend on a second GET call.
# https://storyboard.openstack.org/#!/story/1619675
checksum = resp.headers.get("Content-MD5")
if checksum is None:
# If we don't receive the Content-MD5 header with the download,
# make an additional call to get the image details and look at
# the checksum attribute.
details = self.fetch(session)
checksum = details.checksum
if output:
try:
# In python 2 we might get StringIO - delete it as soon as
# py2 support is dropped
if isinstance(output, io.IOBase) \
or isinstance(output, six.StringIO):
for chunk in resp.iter_content(chunk_size=chunk_size):
output.write(chunk)
else:
with open(output, 'wb') as fd:
for chunk in resp.iter_content(
chunk_size=chunk_size):
fd.write(chunk)
return resp
except Exception as e:
raise exceptions.SDKException(
"Unable to download image: %s" % e)
# if we are returning the repsonse object, ensure that it
# has the content-md5 header so that the caller doesn't
# need to jump through the same hoops through which we
# just jumped.
if stream:
resp.headers['content-md5'] = checksum
return resp
if checksum is not None:
digest = hashlib.md5(resp.content).hexdigest()
if digest != checksum:
raise exceptions.InvalidResponse(
"checksum mismatch: %s != %s" % (checksum, digest))
else:
session.log.warn(
"Unable to verify the integrity of image %s" % (self.id))
return resp

View File

@ -9,17 +9,13 @@
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
import hashlib
import io
import six
from openstack import exceptions
from openstack.image import _download
from openstack import resource
from openstack import utils
class Image(resource.Resource, resource.TagMixin):
class Image(resource.Resource, resource.TagMixin, _download.DownloadMixin):
resources_key = 'images'
base_path = '/images'
@ -274,61 +270,6 @@ class Image(resource.Resource, resource.TagMixin):
'method: "web-download"')
session.post(url, json=json)
def download(self, session, stream=False, output=None, chunk_size=1024):
"""Download the data contained in an image"""
# TODO(briancurtin): This method should probably offload the get
# operation into another thread or something of that nature.
url = utils.urljoin(self.base_path, self.id, 'file')
resp = session.get(url, stream=stream)
# See the following bug report for details on why the checksum
# code may sometimes depend on a second GET call.
# https://storyboard.openstack.org/#!/story/1619675
checksum = resp.headers.get('Content-MD5')
if checksum is None:
# If we don't receive the Content-MD5 header with the download,
# make an additional call to get the image details and look at
# the checksum attribute.
details = self.fetch(session)
checksum = details.checksum
if output:
try:
# In python 2 we might get StringIO - delete it as soon as
# py2 support is dropped
if isinstance(output, io.IOBase) \
or isinstance(output, six.StringIO):
for chunk in resp.iter_content(chunk_size=chunk_size):
output.write(chunk)
else:
with open(output, 'wb') as fd:
for chunk in resp.iter_content(
chunk_size=chunk_size):
fd.write(chunk)
return resp
except Exception as e:
raise exceptions.SDKException(
'Unable to download image: %s' % e)
# if we are returning the repsonse object, ensure that it
# has the content-md5 header so that the caller doesn't
# need to jump through the same hoops through which we
# just jumped.
if stream:
resp.headers['content-md5'] = checksum
return resp
if checksum is not None:
digest = hashlib.md5(resp.content).hexdigest()
if digest != checksum:
raise exceptions.InvalidResponse(
"checksum mismatch: %s != %s" % (checksum, digest))
else:
session.log.warn(
"Unable to verify the integrity of image %s" % (self.id))
return resp
def _prepare_request(self, requires_id=None, prepend_key=False,
patch=False, base_path=None):
request = super(Image, self)._prepare_request(requires_id=requires_id,

View File

@ -220,7 +220,8 @@ def make_fake_stack_event(
def make_fake_image(
image_id=None, md5=NO_MD5, sha256=NO_SHA256, status='active',
image_name=u'fake_image'):
image_name=u'fake_image',
checksum=u'ee36e35a297980dee1b514de9803ec6d'):
return {
u'image_state': u'available',
u'container_format': u'bare',
@ -242,7 +243,7 @@ def make_fake_image(
u'min_disk': 40,
u'virtual_size': None,
u'name': image_name,
u'checksum': u'ee36e35a297980dee1b514de9803ec6d',
u'checksum': checksum,
u'created_at': u'2016-02-10T05:03:11Z',
u'owner_specified.openstack.md5': NO_MD5,
u'owner_specified.openstack.sha256': NO_SHA256,

View File

@ -11,7 +11,7 @@
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
import hashlib
import operator
import tempfile
import uuid
@ -37,10 +37,11 @@ class BaseTestImage(base.TestCase):
self.imagefile = tempfile.NamedTemporaryFile(delete=False)
self.imagefile.write(b'\0')
self.imagefile.close()
self.fake_image_dict = fakes.make_fake_image(
image_id=self.image_id, image_name=self.image_name)
self.fake_search_return = {'images': [self.fake_image_dict]}
self.output = uuid.uuid4().bytes
self.fake_image_dict = fakes.make_fake_image(
image_id=self.image_id, image_name=self.image_name,
checksum=hashlib.md5(self.output).hexdigest())
self.fake_search_return = {'images': [self.fake_image_dict]}
self.container_name = self.getUniqueString('container')

View File

@ -9,7 +9,7 @@
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
import hashlib
import operator
import six
import tempfile
@ -88,6 +88,13 @@ EXAMPLE = {
}
def calculate_md5_checksum(data):
checksum = hashlib.md5()
for chunk in data:
checksum.update(chunk)
return checksum.hexdigest()
class FakeResponse(object):
def __init__(self, response, status_code=200, headers=None, reason=None):
self.body = response
@ -336,8 +343,11 @@ class TestImage(base.TestCase):
self.assertEqual(len(log.records), 1,
"Too many warnings were logged")
self.assertEqual(
"Unable to verify the integrity of image IDENTIFIER",
"Unable to verify the integrity of image %s",
log.records[0].msg)
self.assertEqual(
(sot.id,),
log.records[0].args)
self.sess.get.assert_has_calls(
[mock.call('images/IDENTIFIER/file',
@ -366,6 +376,10 @@ class TestImage(base.TestCase):
response = mock.Mock()
response.status_code = 200
response.iter_content.return_value = [b'01', b'02']
response.headers = {
'Content-MD5':
calculate_md5_checksum(response.iter_content.return_value)
}
self.sess.get = mock.Mock(return_value=response)
sot.download(self.sess, output=output_file)
output_file.seek(0)
@ -376,6 +390,10 @@ class TestImage(base.TestCase):
response = mock.Mock()
response.status_code = 200
response.iter_content.return_value = [b'01', b'02']
response.headers = {
'Content-MD5':
calculate_md5_checksum(response.iter_content.return_value)
}
self.sess.get = mock.Mock(return_value=response)
output_file = tempfile.NamedTemporaryFile()