This commit is contained in:
Pierre Geeraert
@@ -0,0 +1,279 @@
|
||||
__author__ = "John Hollowell"
|
||||
__copyright__ = "(c) John Hollowell 2023"
|
||||
__license__ = "MIT"
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from urllib.parse import urljoin, urlparse
|
||||
|
||||
from proxmoxer import ProxmoxResource, ResourceException
|
||||
from proxmoxer.tools.tasks import Tasks
|
||||
|
||||
CHECKSUM_CHUNK_SIZE = 16384 # read 16k at a time while calculating the checksum for upload
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel(level=logging.WARNING)
|
||||
|
||||
try:
|
||||
import requests
|
||||
except ImportError:
|
||||
logger.error("Files tools requires 'requests' module\n")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
class ChecksumInfo:
|
||||
def __init__(self, name: str, hex_size: int):
|
||||
self.name = name
|
||||
self.hex_size = hex_size
|
||||
|
||||
def __str__(self):
|
||||
return self.name
|
||||
|
||||
def __repr__(self):
|
||||
return f"{self.name} ({self.hex_size} digits)"
|
||||
|
||||
|
||||
class SupportedChecksums(Enum):
|
||||
"""
|
||||
An Enum of the checksum types supported by Proxmox
|
||||
"""
|
||||
|
||||
# ordered by preference for longer/stronger checksums first
|
||||
SHA512 = ChecksumInfo("sha512", 128)
|
||||
SHA256 = ChecksumInfo("sha256", 64)
|
||||
SHA224 = ChecksumInfo("sha224", 56)
|
||||
SHA384 = ChecksumInfo("sha384", 96)
|
||||
MD5 = ChecksumInfo("md5", 32)
|
||||
SHA1 = ChecksumInfo("sha1", 40)
|
||||
|
||||
|
||||
class Files:
|
||||
"""
|
||||
Ease-of-use tools for interacting with the uploading/downloading files
|
||||
in Proxmox VE
|
||||
"""
|
||||
|
||||
def __init__(self, prox: ProxmoxResource, node: str, storage: str):
|
||||
self._prox = prox
|
||||
self._node = node
|
||||
self._storage = storage
|
||||
|
||||
def __repr__(self):
|
||||
return f"Files ({self._node}/{self._storage} at {self._prox})"
|
||||
|
||||
def upload_local_file_to_storage(
|
||||
self,
|
||||
filename: str,
|
||||
do_checksum_check: bool = True,
|
||||
blocking_status: bool = True,
|
||||
):
|
||||
file_path = Path(filename)
|
||||
|
||||
if not file_path.is_file():
|
||||
logger.error(f'"{file_path.absolute()}" does not exist or is not a file')
|
||||
return None
|
||||
|
||||
# init to None in case errors cause no values to be set
|
||||
upid: str = ""
|
||||
checksum: str = None
|
||||
checksum_type: str = None
|
||||
|
||||
try:
|
||||
with open(file_path.absolute(), "rb") as f_obj:
|
||||
if do_checksum_check:
|
||||
# iterate through SupportedChecksums and find the first one in hashlib.algorithms_available
|
||||
for checksum_info in (v.value for v in SupportedChecksums):
|
||||
if checksum_info.name in hashlib.algorithms_available:
|
||||
checksum_type = checksum_info.name
|
||||
break
|
||||
|
||||
if checksum_type is None:
|
||||
logger.warning(
|
||||
"There are no Proxmox supported checksums which are supported by hashlib. Skipping checksum validation"
|
||||
)
|
||||
else:
|
||||
h = hashlib.new(checksum_type)
|
||||
|
||||
# Iterate through the file in CHECKSUM_CHUNK_SIZE size
|
||||
for byte_block in iter(lambda: f_obj.read(CHECKSUM_CHUNK_SIZE), b""):
|
||||
h.update(byte_block)
|
||||
checksum = h.hexdigest()
|
||||
logger.debug(
|
||||
f"The {checksum_type} checksum of {file_path.absolute()} is {checksum}"
|
||||
)
|
||||
|
||||
# reset to the start of the file so the upload can use the same file handle
|
||||
f_obj.seek(0)
|
||||
|
||||
params = {
|
||||
"content": "iso" if file_path.absolute().name.endswith("iso") else "vztmpl",
|
||||
"checksum-algorithm": checksum_type,
|
||||
"checksum": checksum,
|
||||
"filename": f_obj,
|
||||
}
|
||||
upid = self._prox.nodes(self._node).storage(self._storage).upload.post(**params)
|
||||
except OSError as e:
|
||||
logger.error(e)
|
||||
return None
|
||||
|
||||
if blocking_status:
|
||||
return Tasks.blocking_status(self._prox, upid)
|
||||
else:
|
||||
return self._prox.nodes(self._node).tasks(upid).status.get()
|
||||
|
||||
def download_file_to_storage(
|
||||
self,
|
||||
url: str,
|
||||
checksum: Optional[str] = None,
|
||||
checksum_type: Optional[str] = None,
|
||||
blocking_status: bool = True,
|
||||
):
|
||||
file_info = self.get_file_info(url)
|
||||
filename = None
|
||||
|
||||
if file_info is not None:
|
||||
filename = file_info.get("filename")
|
||||
|
||||
if checksum is None and checksum_type is None:
|
||||
checksum, checksum_info = self.get_checksums_from_file_url(url, filename)
|
||||
checksum_type = checksum_info.name if checksum_info else None
|
||||
elif checksum is None or checksum_type is None:
|
||||
logger.error(
|
||||
"Must pass both checksum and checksum_type or leave both None for auto-discovery"
|
||||
)
|
||||
return None
|
||||
|
||||
if checksum is None or checksum_type is None:
|
||||
logger.warning("Unable to discover checksum. Will not do checksum validation")
|
||||
|
||||
params = {
|
||||
"checksum-algorithm": checksum_type,
|
||||
"url": url,
|
||||
"checksum": checksum,
|
||||
"content": "iso" if url.endswith("iso") else "vztmpl",
|
||||
"filename": filename,
|
||||
}
|
||||
upid = self._prox.nodes(self._node).storage(self._storage)("download-url").post(**params)
|
||||
|
||||
if blocking_status:
|
||||
return Tasks.blocking_status(self._prox, upid)
|
||||
else:
|
||||
return self._prox.nodes(self._node).tasks(upid).status.get()
|
||||
|
||||
def get_file_info(self, url: str):
|
||||
try:
|
||||
return self._prox.nodes(self._node)("query-url-metadata").get(url=url)
|
||||
|
||||
except ResourceException as e:
|
||||
logger.warning(f"Unable to get information for {url}: {e}")
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def get_checksums_from_file_url(
|
||||
url: str, filename: str = None, preferred_type=SupportedChecksums.SHA512.value
|
||||
):
|
||||
getters_by_quality = [
|
||||
Files._get_checksum_from_sibling_file,
|
||||
Files._get_checksum_from_extension,
|
||||
Files._get_checksum_from_extension_upper,
|
||||
]
|
||||
|
||||
# hacky way to try the preferred_type first while still trying all types with no duplicates
|
||||
all_types_with_priority = list(
|
||||
dict.fromkeys([preferred_type, *(map(lambda t: t.value, SupportedChecksums))])
|
||||
)
|
||||
for c_info in all_types_with_priority:
|
||||
for getter in getters_by_quality:
|
||||
checksum: str = getter(url, c_info, filename)
|
||||
if checksum is not None:
|
||||
logger.info(f"{getter} found {str(c_info)} checksum {checksum}")
|
||||
return (checksum, c_info)
|
||||
else:
|
||||
logger.debug(f"{getter} found no {str(c_info)} checksum")
|
||||
|
||||
return (None, None)
|
||||
|
||||
@staticmethod
|
||||
def _get_checksum_from_sibling_file(
|
||||
url: str, checksum_info: ChecksumInfo, filename: Optional[str] = None
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Uses a checksum file in the same path as the target file to discover the checksum
|
||||
|
||||
:param url: the URL string of the target file
|
||||
:type url: str
|
||||
:param checksum_info: the type of checksum to search for
|
||||
:type checksum_info: ChecksumInfo
|
||||
:param filename: the filename to use for finding the checksum. If None, it will be discovered from the url
|
||||
:type filename: str | None
|
||||
:return: a string of the checksum if found, else None
|
||||
:rtype: str | None
|
||||
"""
|
||||
sumfile_url = urljoin(url, (checksum_info.name + "SUMS").upper())
|
||||
filename = filename or os.path.basename(urlparse(url).path)
|
||||
|
||||
return Files._get_checksum_helper(sumfile_url, filename, checksum_info)
|
||||
|
||||
@staticmethod
|
||||
def _get_checksum_from_extension(
|
||||
url: str, checksum_info: ChecksumInfo, filename: Optional[str] = None
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Uses a checksum file with a checksum extension added to the target file to discover the checksum
|
||||
|
||||
:param url: the URL string of the target file
|
||||
:type url: str
|
||||
:param checksum_info: the type of checksum to search for
|
||||
:type checksum_info: ChecksumInfo
|
||||
:param filename: the filename to use for finding the checksum. If None, it will be discovered from the url
|
||||
:type filename: str | None
|
||||
:return: a string of the checksum if found, else None
|
||||
:rtype: str | None
|
||||
"""
|
||||
sumfile_url = url + "." + checksum_info.name
|
||||
filename = filename or os.path.basename(urlparse(url).path)
|
||||
|
||||
return Files._get_checksum_helper(sumfile_url, filename, checksum_info)
|
||||
|
||||
@staticmethod
|
||||
def _get_checksum_from_extension_upper(
|
||||
url: str, checksum_info: ChecksumInfo, filename: Optional[str] = None
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Uses a checksum file with a checksum extension added to the target file to discover the checksum
|
||||
|
||||
:param url: the URL string of the target file
|
||||
:type url: str
|
||||
:param checksum_info: the type of checksum to search for
|
||||
:type checksum_info: ChecksumInfo
|
||||
:param filename: the filename to use for finding the checksum. If None, it will be discovered from the url
|
||||
:type filename: str | None
|
||||
:return: a string of the checksum if found, else None
|
||||
:rtype: str | None
|
||||
"""
|
||||
sumfile_url = url + "." + checksum_info.name.upper()
|
||||
filename = filename or os.path.basename(urlparse(url).path)
|
||||
|
||||
return Files._get_checksum_helper(sumfile_url, filename, checksum_info)
|
||||
|
||||
@staticmethod
|
||||
def _get_checksum_helper(sumfile_url: str, filename: str, checksum_info: ChecksumInfo):
|
||||
logger.debug(f"getting {sumfile_url}")
|
||||
try:
|
||||
resp = requests.get(sumfile_url, timeout=10)
|
||||
except (requests.exceptions.ConnectionError, requests.exceptions.ReadTimeout):
|
||||
logger.info(f"Failed when trying to get {sumfile_url}")
|
||||
return None
|
||||
|
||||
if resp.status_code == 200:
|
||||
for line in resp.iter_lines():
|
||||
line_str = line.decode("utf-8")
|
||||
logger.debug(f"checking for '{filename}' in '{line_str}'")
|
||||
if filename in str(line_str):
|
||||
return line_str[0 : checksum_info.hex_size]
|
||||
return None
|
||||
Reference in New Issue
Block a user