Add safe ISO download and explicit run/open commands

tests/test_media.py

@@ -0,0 +1,134 @@
+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'<html>download here</html>' 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('<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'])