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:
parent
1e810595c6
commit
cf9922e885
|
@ -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
|
|
@ -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
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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')
|
||||
|
||||
|
||||
|
|
|
@ -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()
|
||||
|
|
Loading…
Reference in New Issue