211 lines
10 KiB
Python
211 lines
10 KiB
Python
import hashlib
|
|
import http.server
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from hydra import media
|
|
|
|
|
|
def iso_bytes(signature=b'CD001', *, declared_blocks=40, block_size=2048, hybrid=False):
|
|
data = bytearray(40 * 2048)
|
|
start = 16 * 2048
|
|
data[start:start + 7] = b'\x01' + signature + b'\x01'
|
|
data[start + 80:start + 88] = declared_blocks.to_bytes(4, 'little') + declared_blocks.to_bytes(4, 'big')
|
|
data[start + 120:start + 124] = b'\x01\x00\x00\x01'
|
|
data[start + 124:start + 128] = b'\x01\x00\x00\x01'
|
|
data[start + 128:start + 132] = block_size.to_bytes(2, 'little') + block_size.to_bytes(2, 'big')
|
|
data[start + 881] = 1
|
|
data[17 * 2048:17 * 2048 + 7] = b'\xffCD001\x01'
|
|
if hybrid:
|
|
data[20 * 2048:20 * 2048 + 7] = b'\x00NSR03\x01'
|
|
return bytes(data)
|
|
|
|
|
|
class MediaTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.tmp.cleanup)
|
|
self.cache = Path(self.tmp.name) / 'cache'
|
|
self.requests = []
|
|
owner = self
|
|
class Handler(http.server.BaseHTTPRequestHandler):
|
|
def do_GET(self):
|
|
owner.requests.append(self.path)
|
|
body = b'<html>download here</html>' if self.path.startswith('/html') else iso_bytes(
|
|
declared_blocks=4000 if self.path.startswith('/extent') else 40)
|
|
self.send_response(200)
|
|
self.send_header('Content-Type', 'text/html' if self.path.startswith('/html') else 'application/octet-stream')
|
|
if not self.path.startswith('/extent-no-length'):
|
|
self.send_header('Content-Length', str(len(body) + (100 if self.path.startswith('/truncated') else 0)))
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
def log_message(self, *args):
|
|
pass
|
|
self.server = http.server.ThreadingHTTPServer(('127.0.0.1', 0), Handler)
|
|
self.addCleanup(self.server.server_close)
|
|
self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
|
|
self.thread.start()
|
|
self.addCleanup(self.server.shutdown)
|
|
self.url = f'http://127.0.0.1:{self.server.server_port}'
|
|
|
|
def test_verified_download_is_reused_and_progress_reports_complete_bytes(self):
|
|
updates = []
|
|
url = self.url + '/installer.iso?token=secret'
|
|
path = media.download_iso(url, cache=self.cache, progress=lambda got, total: updates.append((got, total)))
|
|
self.assertEqual(path.read_bytes(), iso_bytes())
|
|
self.assertEqual(path.stat().st_mode & 0o777, 0o600)
|
|
self.assertNotIn('secret', path.name)
|
|
self.assertEqual(updates[-1], (81920, 81920))
|
|
self.assertEqual(media.download_iso(url, cache=self.cache), path)
|
|
self.assertEqual(len(self.requests), 1)
|
|
|
|
def test_html_truncation_and_checksum_do_not_publish_images(self):
|
|
for endpoint, checksum in (('/html.iso', None), ('/truncated.iso', None), ('/wrong.iso', '0' * 64)):
|
|
with self.subTest(endpoint=endpoint):
|
|
with self.assertRaises(ValueError):
|
|
media.download_iso(self.url + endpoint, cache=self.cache, sha256=checksum)
|
|
self.assertFalse(list(self.cache.glob('*.iso')))
|
|
self.assertFalse(list(self.cache.glob('*.part')))
|
|
|
|
def test_cancellation_removes_partial_and_retry_completes(self):
|
|
def cancel(got, total):
|
|
if got:
|
|
raise KeyboardInterrupt
|
|
with self.assertRaises(KeyboardInterrupt):
|
|
media.download_iso(self.url + '/cancel.iso', cache=self.cache, progress=cancel)
|
|
self.assertFalse(list(self.cache.glob('*.iso')))
|
|
self.assertFalse(list(self.cache.glob('*.part')))
|
|
self.assertEqual(media.download_iso(self.url + '/cancel.iso', cache=self.cache).read_bytes(), iso_bytes())
|
|
|
|
def test_same_basename_different_urls_and_concurrent_downloads_are_safe(self):
|
|
results = []
|
|
errors = []
|
|
def fetch():
|
|
try:
|
|
results.append(media.download_iso(self.url + '/one/os.iso', cache=self.cache))
|
|
except Exception as exc:
|
|
errors.append(exc)
|
|
threads = [threading.Thread(target=fetch) for _ in range(2)]
|
|
for thread in threads:
|
|
thread.start()
|
|
for thread in threads:
|
|
thread.join(5)
|
|
self.assertEqual(errors, [])
|
|
self.assertEqual(len(set(results)), 1)
|
|
self.assertEqual(self.requests, ['/one/os.iso'])
|
|
second = media.download_iso(self.url + '/two/os.iso', cache=self.cache)
|
|
self.assertNotEqual(second, results[0])
|
|
|
|
def test_corrupt_cache_and_symlink_never_overwrite_unrelated_files(self):
|
|
url = self.url + '/installer.iso'
|
|
path = media.download_iso(url, cache=self.cache)
|
|
path.write_text('valuable unrelated data')
|
|
with self.assertRaises(ValueError):
|
|
media.download_iso(url, cache=self.cache)
|
|
self.assertEqual(path.read_text(), 'valuable unrelated data')
|
|
path.unlink()
|
|
target = Path(self.tmp.name) / 'outside'
|
|
target.write_bytes(iso_bytes())
|
|
path.symlink_to(target)
|
|
with self.assertRaises(ValueError):
|
|
media.download_iso(url, cache=self.cache)
|
|
self.assertTrue(path.is_symlink())
|
|
self.assertEqual(target.read_bytes(), iso_bytes())
|
|
|
|
def test_unsupported_urls_checksums_and_insufficient_space_fail_before_download(self):
|
|
for url in ('ftp://example.org/os.iso', 'http://user:secret@example.org/os.iso'):
|
|
with self.assertRaises(ValueError):
|
|
media.download_iso(url, cache=self.cache)
|
|
with self.assertRaises(ValueError):
|
|
media.download_iso(self.url + '/os.iso', cache=self.cache, sha256='bad')
|
|
with patch('hydra.media.shutil.disk_usage', return_value=(10, 9, 1)):
|
|
with self.assertRaisesRegex(ValueError, 'space'):
|
|
media.download_iso(self.url + '/os.iso', cache=self.cache)
|
|
self.assertFalse(list(self.cache.glob('*.iso')))
|
|
|
|
def test_local_validation_accepts_iso_udf_hybrid_and_rejects_html_and_bad_checksum(self):
|
|
path = Path(self.tmp.name) / 'os.iso'
|
|
path.write_bytes(iso_bytes(hybrid=True))
|
|
media.validate_iso(path, sha256=hashlib.sha256(path.read_bytes()).hexdigest())
|
|
with self.assertRaises(ValueError):
|
|
media.validate_iso(path, sha256='0' * 64)
|
|
path.write_text('<html>not an image</html>')
|
|
with self.assertRaisesRegex(ValueError, 'direct|ISO'):
|
|
media.validate_iso(path)
|
|
|
|
def test_cache_reuse_reports_completed_progress_without_another_request(self):
|
|
url = self.url + '/cached.iso'
|
|
path = media.download_iso(url, cache=self.cache)
|
|
updates = []
|
|
self.assertEqual(media.download_iso(url, cache=self.cache, progress=lambda got, total: updates.append((got, total))), path)
|
|
self.assertEqual(updates, [(81920, 81920)])
|
|
self.assertEqual(self.requests, ['/cached.iso'])
|
|
|
|
def test_declared_extent_truncation_with_matching_content_length_never_publishes(self):
|
|
with self.assertRaisesRegex(media.MediaError, 'truncated|extent'):
|
|
media.download_iso(self.url + '/extent.iso', cache=self.cache)
|
|
self.assertFalse(list(self.cache.glob('*.iso')))
|
|
self.assertFalse(list(self.cache.glob('*.part')))
|
|
|
|
def test_declared_extent_truncation_without_content_length_never_publishes(self):
|
|
with self.assertRaisesRegex(media.MediaError, 'truncated|extent'):
|
|
media.download_iso(self.url + '/extent-no-length.iso', cache=self.cache)
|
|
self.assertFalse(list(self.cache.glob('*.iso')))
|
|
self.assertFalse(list(self.cache.glob('*.part')))
|
|
|
|
def test_zero_invalid_or_inconsistent_block_fields_are_rejected(self):
|
|
path = Path(self.tmp.name) / 'invalid.iso'
|
|
for body in (iso_bytes(declared_blocks=0), iso_bytes(block_size=0), iso_bytes(block_size=4096)):
|
|
with self.subTest(body=body[16 * 2048 + 80:16 * 2048 + 132]):
|
|
path.write_bytes(body)
|
|
with self.assertRaises(media.MediaError):
|
|
media.validate_iso(path)
|
|
body = bytearray(iso_bytes())
|
|
body[16 * 2048 + 84:16 * 2048 + 88] = (4000).to_bytes(4, 'big')
|
|
path.write_bytes(body)
|
|
with self.assertRaises(media.MediaError):
|
|
media.validate_iso(path)
|
|
body = bytearray(iso_bytes())
|
|
body[16 * 2048 + 130:16 * 2048 + 132] = (1024).to_bytes(2, 'big')
|
|
path.write_bytes(body)
|
|
with self.assertRaises(media.MediaError):
|
|
media.validate_iso(path)
|
|
|
|
def test_boot_signature_or_unterminated_primary_is_not_a_complete_image(self):
|
|
path = Path(self.tmp.name) / 'descriptor.iso'
|
|
boot_only = bytearray(iso_bytes())
|
|
boot_only[16 * 2048] = 0
|
|
unterminated = bytearray(iso_bytes())
|
|
unterminated[17 * 2048:18 * 2048] = bytes(2048)
|
|
for body in (boot_only, unterminated, iso_bytes()[:17 * 2048]):
|
|
path.write_bytes(body)
|
|
with self.assertRaises(media.MediaError):
|
|
media.validate_iso(path)
|
|
|
|
def test_pure_udf_is_rejected_with_supported_image_guidance(self):
|
|
path = Path(self.tmp.name) / 'udf.iso'
|
|
for signature in (b'NSR02', b'NSR03'):
|
|
path.write_bytes(iso_bytes(signature))
|
|
with self.assertRaisesRegex(media.MediaError, 'UDF.*ISO9660|ISO9660.*UDF'):
|
|
media.validate_iso(path)
|
|
|
|
def test_iso_udf_hybrid_with_trailing_data_and_smaller_blocks_remains_valid(self):
|
|
path = Path(self.tmp.name) / 'hybrid.iso'
|
|
for body in (iso_bytes(hybrid=True) + b'trailing hybrid partition data',
|
|
iso_bytes(declared_blocks=80, block_size=1024),
|
|
iso_bytes(declared_blocks=160, block_size=512)):
|
|
path.write_bytes(body)
|
|
self.assertEqual(media.validate_iso(path), path)
|
|
|
|
def test_cached_iso_truncated_with_descriptor_intact_is_not_reused(self):
|
|
url = self.url + '/cached-truncated.iso'
|
|
path = media.download_iso(url, cache=self.cache)
|
|
path.write_bytes(path.read_bytes()[:20 * 2048])
|
|
with self.assertRaisesRegex(media.MediaError, 'Cached ISO failed validation'):
|
|
media.download_iso(url, cache=self.cache)
|
|
self.assertEqual(path.stat().st_size, 20 * 2048)
|
|
self.assertEqual(self.requests, ['/cached-truncated.iso'])
|