diff --git a/openstack/image/_download.py b/openstack/image/_download.py new file mode 100644 index 000000000..83130e8bb --- /dev/null +++ b/openstack/image/_download.py @@ -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 diff --git a/openstack/image/v1/image.py b/openstack/image/v1/image.py index 6760ed51f..a92c4362d 100644 --- a/openstack/image/v1/image.py +++ b/openstack/image/v1/image.py @@ -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 diff --git a/openstack/image/v2/image.py b/openstack/image/v2/image.py index cc6a3236c..50af9e8cf 100644 --- a/openstack/image/v2/image.py +++ b/openstack/image/v2/image.py @@ -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, diff --git a/openstack/tests/fakes.py b/openstack/tests/fakes.py index 9055075b4..4479968c3 100644 --- a/openstack/tests/fakes.py +++ b/openstack/tests/fakes.py @@ -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, diff --git a/openstack/tests/unit/cloud/test_image.py b/openstack/tests/unit/cloud/test_image.py index f972f7603..c90a28f5c 100644 --- a/openstack/tests/unit/cloud/test_image.py +++ b/openstack/tests/unit/cloud/test_image.py @@ -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') diff --git a/openstack/tests/unit/image/v2/test_image.py b/openstack/tests/unit/image/v2/test_image.py index 281c89f4c..7c429e48d 100644 --- a/openstack/tests/unit/image/v2/test_image.py +++ b/openstack/tests/unit/image/v2/test_image.py @@ -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()