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'): data = bytearray(40 * 2048) data[16 * 2048:16 * 2048 + 7] = b'\x01' + signature + b'\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'download here' if self.path.startswith('/html') else iso_bytes() self.send_response(200) self.send_header('Content-Type', 'text/html' if self.path.startswith('/html') else 'application/octet-stream') 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_udf_and_rejects_html_and_bad_checksum(self): path = Path(self.tmp.name) / 'os.iso' path.write_bytes(iso_bytes(b'NSR03')) 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('not an image') 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'])