336 lines
9.2 KiB
Python
336 lines
9.2 KiB
Python
#
|
|
# Copyright 2006-2007, 2013 Red Hat, Inc.
|
|
# Daniel P. Berrange <berrange@redhat.com>
|
|
#
|
|
# This work is licensed under the GNU GPLv2 or later.
|
|
# See the COPYING file in the top-level directory.
|
|
|
|
import ftplib
|
|
import io
|
|
import logging
|
|
import os
|
|
import subprocess
|
|
import tempfile
|
|
import urllib
|
|
|
|
import requests
|
|
|
|
|
|
###########################################################################
|
|
# Backends for the various URL types we support (http, https, ftp, local) #
|
|
###########################################################################
|
|
|
|
class _URLFetcher(object):
|
|
"""
|
|
This is a generic base class for fetching/extracting files from
|
|
a media source, such as CD ISO, or HTTP/HTTPS/FTP server
|
|
"""
|
|
_block_size = 16384
|
|
|
|
def __init__(self, location, scratchdir, meter):
|
|
self.location = location
|
|
self.scratchdir = scratchdir
|
|
self.meter = meter
|
|
|
|
self._srcdir = None
|
|
|
|
logging.debug("Using scratchdir=%s", scratchdir)
|
|
|
|
|
|
####################
|
|
# Internal helpers #
|
|
####################
|
|
|
|
def _make_full_url(self, filename):
|
|
"""
|
|
Generate a full fetchable URL from the passed filename, which
|
|
is relative to the self.location
|
|
"""
|
|
ret = self._srcdir or self.location
|
|
if not filename:
|
|
return ret
|
|
|
|
if not ret.endswith("/"):
|
|
ret += "/"
|
|
return ret + filename
|
|
|
|
def _grabURL(self, filename, fileobj):
|
|
"""
|
|
Download the filename from self.location, and write contents to
|
|
fileobj
|
|
"""
|
|
url = self._make_full_url(filename)
|
|
|
|
try:
|
|
urlobj, size = self._grabber(url)
|
|
except Exception as e:
|
|
raise ValueError(_("Couldn't acquire file %s: %s") %
|
|
(url, str(e)))
|
|
|
|
logging.debug("Fetching URI: %s", url)
|
|
self.meter.start(
|
|
text=_("Retrieving file %s...") % os.path.basename(filename),
|
|
size=size)
|
|
|
|
total = self._write(urlobj, fileobj)
|
|
self.meter.end(total)
|
|
|
|
def _write(self, urlobj, fileobj):
|
|
"""
|
|
Write the contents of urlobj to python file like object fileobj
|
|
"""
|
|
total = 0
|
|
while 1:
|
|
buff = urlobj.read(self._block_size)
|
|
if not buff:
|
|
break
|
|
fileobj.write(buff)
|
|
total += len(buff)
|
|
self.meter.update(total)
|
|
return total
|
|
|
|
def _grabber(self, url):
|
|
"""
|
|
Returns the urlobj, size for the passed URL. urlobj is whatever
|
|
data needs to be passed to self._write
|
|
"""
|
|
raise NotImplementedError("must be implemented in subclass")
|
|
|
|
|
|
##############
|
|
# Public API #
|
|
##############
|
|
|
|
def prepareLocation(self):
|
|
"""
|
|
Perform any necessary setup
|
|
"""
|
|
pass
|
|
|
|
def cleanupLocation(self):
|
|
"""
|
|
Perform any necessary cleanup
|
|
"""
|
|
pass
|
|
|
|
def can_access(self):
|
|
"""
|
|
Return True if the location URL seems to be valid
|
|
"""
|
|
return True
|
|
|
|
def _hasFile(self, url):
|
|
raise NotImplementedError("Must be implemented in subclass")
|
|
|
|
def hasFile(self, filename):
|
|
"""
|
|
Return True if self.location has the passed filename
|
|
"""
|
|
url = self._make_full_url(filename)
|
|
ret = self._hasFile(url)
|
|
logging.debug("hasFile(%s) returning %s", url, ret)
|
|
return ret
|
|
|
|
def acquireFile(self, filename):
|
|
"""
|
|
Grab the passed filename from self.location and save it to
|
|
a temporary file, returning the temp filename
|
|
"""
|
|
prefix = "virtinst-" + os.path.basename(filename) + "."
|
|
|
|
# pylint: disable=redefined-variable-type
|
|
if "VIRTINST_TEST_SUITE" in os.environ:
|
|
fn = os.path.join("/tmp", prefix)
|
|
fileobj = open(fn, "wb")
|
|
else:
|
|
fileobj = tempfile.NamedTemporaryFile(
|
|
dir=self.scratchdir, prefix=prefix, delete=False)
|
|
fn = fileobj.name
|
|
|
|
self._grabURL(filename, fileobj)
|
|
logging.debug("Saved file to %s", fn)
|
|
return fn
|
|
|
|
def acquireFileContent(self, filename):
|
|
"""
|
|
Grab the passed filename from self.location and return it as a string
|
|
"""
|
|
fileobj = io.BytesIO()
|
|
self._grabURL(filename, fileobj)
|
|
return fileobj.getvalue().decode("utf-8")
|
|
|
|
|
|
class _HTTPURLFetcher(_URLFetcher):
|
|
_session = None
|
|
|
|
def prepareLocation(self):
|
|
self._session = requests.Session()
|
|
|
|
def cleanupLocation(self):
|
|
if self._session:
|
|
try:
|
|
self._session.close()
|
|
except Exception:
|
|
logging.debug("Error closing requests.session", exc_info=True)
|
|
self._session = None
|
|
|
|
def can_access(self):
|
|
return self.hasFile("")
|
|
|
|
def _hasFile(self, url):
|
|
"""
|
|
We just do a HEAD request to see if the file exists
|
|
"""
|
|
try:
|
|
response = self._session.head(url, allow_redirects=True)
|
|
response.raise_for_status()
|
|
except Exception as e:
|
|
logging.debug("HTTP hasFile request failed: %s", str(e))
|
|
return False
|
|
return True
|
|
|
|
def _grabber(self, url):
|
|
"""
|
|
Use requests for this
|
|
"""
|
|
response = self._session.get(url, stream=True)
|
|
response.raise_for_status()
|
|
try:
|
|
size = int(response.headers.get('content-length'))
|
|
except Exception:
|
|
size = None
|
|
return response, size
|
|
|
|
def _write(self, urlobj, fileobj):
|
|
"""
|
|
The requests object doesn't have a file-like read() option, so
|
|
we need to implement it ourselves
|
|
"""
|
|
total = 0
|
|
for data in urlobj.iter_content(chunk_size=self._block_size):
|
|
fileobj.write(data)
|
|
total += len(data)
|
|
self.meter.update(total)
|
|
return total
|
|
|
|
|
|
class _FTPURLFetcher(_URLFetcher):
|
|
_ftp = None
|
|
|
|
def prepareLocation(self):
|
|
if self._ftp:
|
|
return
|
|
|
|
try:
|
|
parsed = urllib.parse.urlparse(self.location)
|
|
self._ftp = ftplib.FTP()
|
|
username = urllib.parse.unquote(parsed.username or '')
|
|
password = urllib.parse.unquote(parsed.password or '')
|
|
self._ftp.connect(parsed.hostname, parsed.port or 0)
|
|
self._ftp.login(username, password)
|
|
# Force binary mode
|
|
self._ftp.voidcmd("TYPE I")
|
|
except Exception as e:
|
|
raise ValueError(_("Opening URL %s failed: %s.") %
|
|
(self.location, str(e)))
|
|
|
|
def _grabber(self, url):
|
|
"""
|
|
Use urllib and ftplib to grab the file
|
|
"""
|
|
request = urllib.request.Request(url)
|
|
urlobj = urllib.request.urlopen(request)
|
|
size = self._ftp.size(urllib.parse.urlparse(url)[2])
|
|
return urlobj, size
|
|
|
|
|
|
def cleanupLocation(self):
|
|
if not self._ftp:
|
|
return
|
|
|
|
try:
|
|
self._ftp.quit()
|
|
except Exception:
|
|
logging.debug("Error quitting ftp connection", exc_info=True)
|
|
|
|
self._ftp = None
|
|
|
|
def _hasFile(self, url):
|
|
path = urllib.parse.urlparse(url)[2]
|
|
|
|
try:
|
|
try:
|
|
# If it's a file
|
|
self._ftp.size(path)
|
|
except ftplib.all_errors:
|
|
# If it's a dir
|
|
self._ftp.cwd(path)
|
|
except ftplib.all_errors as e:
|
|
logging.debug("FTP hasFile: couldn't access %s: %s",
|
|
url, str(e))
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
class _LocalURLFetcher(_URLFetcher):
|
|
"""
|
|
For grabbing files from a local directory
|
|
"""
|
|
def _hasFile(self, url):
|
|
return os.path.exists(url)
|
|
|
|
def _grabber(self, url):
|
|
urlobj = open(url, "rb")
|
|
size = os.path.getsize(url)
|
|
return urlobj, size
|
|
|
|
|
|
class _ISOURLFetcher(_URLFetcher):
|
|
_cache_file_list = None
|
|
|
|
def _make_full_url(self, filename):
|
|
return "/" + filename
|
|
|
|
def _grabber(self, url):
|
|
"""
|
|
Use isoinfo to grab the file
|
|
"""
|
|
if not self._hasFile(url):
|
|
raise RuntimeError("isoinfo didn't find file=%s" % url)
|
|
|
|
cmd = ["isoinfo", "-J", "-i", self.location, "-x", url]
|
|
|
|
logging.debug("Running isoinfo: %s", cmd)
|
|
output = subprocess.check_output(cmd)
|
|
|
|
return io.BytesIO(output), len(output)
|
|
|
|
def _hasFile(self, url):
|
|
"""
|
|
Use isoinfo to list and search for the file
|
|
"""
|
|
if not self._cache_file_list:
|
|
cmd = ["isoinfo", "-J", "-i", self.location, "-f"]
|
|
|
|
logging.debug("Running isoinfo: %s", cmd)
|
|
output = subprocess.check_output(cmd)
|
|
|
|
self._cache_file_list = output.splitlines(False)
|
|
|
|
return url.encode("ascii") in self._cache_file_list
|
|
|
|
|
|
def fetcherForURI(uri, *args, **kwargs):
|
|
if uri.startswith("http://") or uri.startswith("https://"):
|
|
fclass = _HTTPURLFetcher
|
|
elif uri.startswith("ftp://"):
|
|
fclass = _FTPURLFetcher
|
|
elif os.path.isdir(uri):
|
|
# Pointing to a local tree
|
|
fclass = _LocalURLFetcher
|
|
else:
|
|
# Pointing to a path (e.g. iso), or a block device (e.g. /dev/cdrom)
|
|
fclass = _ISOURLFetcher
|
|
return fclass(uri, *args, **kwargs)
|