9d6211f1b3
Python will eventually close open files, but it is recommended to explicitly close them instead of waiting for the gc or program exit. This fixes all the uses of open...read/write in the codebase, mostly in tests.
360 lines
18 KiB
Python
360 lines
18 KiB
Python
#
|
|
# Copyright (C) 2018 Red Hat, Inc.
|
|
#
|
|
# This program is free software; you can redistribute it and/or modify
|
|
# it under the terms of the GNU General Public License as published by
|
|
# the Free Software Foundation; either version 2 of the License, or
|
|
# (at your option) any later version.
|
|
#
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU General Public License for more details.
|
|
#
|
|
# You should have received a copy of the GNU General Public License
|
|
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
|
#
|
|
import glob
|
|
import os
|
|
import parted
|
|
from subprocess import CalledProcessError
|
|
import tarfile
|
|
import tempfile
|
|
import unittest
|
|
|
|
from ..lib import get_file_magic
|
|
from pylorax.executils import runcmd
|
|
from pylorax.imgutils import mkcpio, mktar, mksquashfs, mksparse, mkqcow2, loop_attach, loop_detach
|
|
from pylorax.imgutils import get_loop_name, LoopDev, dm_attach, dm_detach, DMDev, Mount
|
|
from pylorax.imgutils import mkdosimg, mkext4img, mkbtrfsimg, mkhfsimg, default_image_name
|
|
from pylorax.imgutils import mount, umount, kpartx_disk_img, PartitionMount, mkfsimage_from_disk
|
|
from pylorax.sysutils import joinpaths
|
|
|
|
def mkfakerootdir(rootdir):
|
|
"""Populate a fake rootdir with a few directories and files
|
|
|
|
:param rootdir: An existing directory to create files/dirs under
|
|
:type rootdir: str
|
|
|
|
Use this for testing the mk* functions that compress a directory tree
|
|
"""
|
|
dirs = ["/root", "/usr/sbin/", "/usr/local/", "/home/bart", "/etc/"]
|
|
files = ["/etc/passwd", "/home/bart/.bashrc", "/root/.bashrc"]
|
|
for d in dirs:
|
|
os.makedirs(joinpaths(rootdir, d))
|
|
for f in files:
|
|
if not os.path.isdir(joinpaths(rootdir, os.path.dirname(f))):
|
|
os.makedirs(joinpaths(rootdir, os.path.dirname(f)))
|
|
with open(joinpaths(rootdir, f), "w") as ff:
|
|
ff.write("I AM FAKE FILE %s" % f.upper())
|
|
|
|
def mkfakebootdir(bootdir):
|
|
"""Populate a fake /boot directory with a kernel and initrd
|
|
|
|
:param bootdir: An existing directory to create files/dirs under
|
|
:type bootdir: str
|
|
"""
|
|
with open(joinpaths(bootdir, "vmlinuz-4.18.13-200.fc28.x86_64"), "w") as f:
|
|
f.write("I AM A FAKE KERNEL")
|
|
with open(joinpaths(bootdir, "initramfs-4.18.13-200.fc28.x86_64.img"), "w") as f:
|
|
f.write("I AM A FAKE INITRD")
|
|
|
|
def mkfakediskimg(disk_img):
|
|
"""Create a fake partitioned disk image
|
|
|
|
:param disk_img: Full path to a partitioned disk image
|
|
:type disk_img: str
|
|
:returns: True if it was successful, False if something went wrong
|
|
|
|
Include /boot, swap, and / partitions with fake kernel and /etc/passwd
|
|
"""
|
|
try:
|
|
mksparse(disk_img, 42 * 1024**2)
|
|
# Make a /boot, / and swap partitions on it
|
|
dev = parted.getDevice(disk_img)
|
|
disk = parted.freshDisk(dev, "gpt")
|
|
|
|
# (start, length, flags, name)
|
|
for start, length, flags, name in [
|
|
( 1024**2, 1024**2, None, "boot"),
|
|
(2*1024**2, 2*1024**2, parted.PARTITION_SWAP, "swap"),
|
|
(4*1024**2, 38*1024**2, None, "root")]:
|
|
geo = parted.Geometry(device=dev, start=start//dev.sectorSize, length=length//dev.sectorSize)
|
|
part = parted.Partition(disk=disk, type=parted.PARTITION_NORMAL, geometry=geo)
|
|
part.getPedPartition().set_name(name)
|
|
disk.addPartition(partition=part)
|
|
if flags:
|
|
part.setFlag(flags)
|
|
disk.commit()
|
|
os.sync()
|
|
except parted.PartedException:
|
|
return False
|
|
|
|
# Mount the disk's partitions
|
|
loop_devs = kpartx_disk_img(disk_img)
|
|
|
|
try:
|
|
# Format the partitions
|
|
runcmd(["mkfs.ext4", "/dev/mapper/" + loop_devs[0][0]])
|
|
runcmd(["mkswap", "/dev/mapper/" + loop_devs[1][0]])
|
|
runcmd(["mkfs.ext4", "/dev/mapper/" + loop_devs[2][0]])
|
|
|
|
# Mount the boot partition and make a fake kernel and initrd
|
|
boot_mnt = mount("/dev/mapper/" + loop_devs[0][0])
|
|
try:
|
|
mkfakebootdir(boot_mnt)
|
|
finally:
|
|
umount(boot_mnt)
|
|
|
|
# Mount the / partition and make a fake / filesystem with /etc/passwd
|
|
root_mnt = mount("/dev/mapper/" + loop_devs[2][0])
|
|
try:
|
|
mkfakerootdir(root_mnt)
|
|
finally:
|
|
umount(root_mnt)
|
|
except Exception:
|
|
return False
|
|
finally:
|
|
# Remove the disk's mounted partitions
|
|
runcmd(["kpartx", "-d", "-s", disk_img])
|
|
|
|
return True
|
|
|
|
class ImgUtilsTest(unittest.TestCase):
|
|
def test_mkcpio(self):
|
|
"""Test mkcpio function"""
|
|
with tempfile.TemporaryDirectory(prefix="lorax.test.") as work_dir:
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img:
|
|
mkfakerootdir(work_dir)
|
|
mkcpio(work_dir, disk_img.name, compression=None)
|
|
|
|
self.assertTrue(os.path.exists(disk_img.name))
|
|
file_details = get_file_magic(disk_img.name)
|
|
self.assertTrue("cpio" in file_details, file_details)
|
|
|
|
def test_mktar(self):
|
|
"""Test mktar function"""
|
|
with tempfile.TemporaryDirectory(prefix="lorax.test.") as work_dir:
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img:
|
|
mkfakerootdir(work_dir)
|
|
mktar(work_dir, disk_img.name, compression=None)
|
|
|
|
self.assertTrue(os.path.exists(disk_img.name))
|
|
file_details = get_file_magic(disk_img.name)
|
|
self.assertTrue("POSIX tar" in file_details, file_details)
|
|
|
|
def test_compressed_mktar(self):
|
|
"""Test compressed mktar function"""
|
|
with tempfile.TemporaryDirectory(prefix="lorax.test.") as work_dir:
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img:
|
|
mkfakerootdir(work_dir)
|
|
for (compression, magic) in [("xz", "XZ compressed"),
|
|
("lzma", "LZMA compressed"),
|
|
("gzip", "gzip compressed"),
|
|
("bzip2", "bzip2 compressed")]:
|
|
os.unlink(disk_img.name)
|
|
mktar(work_dir, disk_img.name, compression=compression)
|
|
|
|
self.assertTrue(os.path.exists(disk_img.name))
|
|
file_details = get_file_magic(disk_img.name)
|
|
self.assertTrue(magic in file_details, (compression, magic, file_details))
|
|
|
|
def test_mktar_single_file(self):
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img,\
|
|
tempfile.NamedTemporaryFile(prefix="lorax.test.input.") as input_file:
|
|
mktar(input_file.name, disk_img.name, compression=None)
|
|
|
|
self.assertTrue(os.path.exists(disk_img.name))
|
|
self.assertTrue(tarfile.is_tarfile(disk_img.name))
|
|
|
|
with tarfile.TarFile(disk_img.name) as t:
|
|
self.assertEqual(t.getnames(), [os.path.basename(input_file.name)])
|
|
|
|
def test_mksquashfs(self):
|
|
"""Test mksquashfs function"""
|
|
with tempfile.TemporaryDirectory(prefix="lorax.test.") as work_dir:
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img:
|
|
mkfakerootdir(work_dir)
|
|
disk_img.close()
|
|
mksquashfs(work_dir, disk_img.name)
|
|
|
|
self.assertTrue(os.path.exists(disk_img.name))
|
|
file_details = get_file_magic(disk_img.name)
|
|
self.assertTrue("Squashfs" in file_details, file_details)
|
|
|
|
def test_mksparse(self):
|
|
"""Test mksparse function"""
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img:
|
|
mksparse(disk_img.name, 42 * 1024**2)
|
|
self.assertEqual(os.stat(disk_img.name).st_size, 42 * 1024**2)
|
|
|
|
def test_mkqcow2(self):
|
|
"""Test mkqcow2 function"""
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img:
|
|
mkqcow2(disk_img.name, 42 * 1024**2)
|
|
file_details = get_file_magic(disk_img.name)
|
|
self.assertTrue("QEMU QCOW" in file_details, file_details)
|
|
self.assertTrue(str(42 * 1024**2) in file_details, file_details)
|
|
|
|
@unittest.skipUnless(os.geteuid() == 0 and not os.path.exists("/.in-container"), "requires root privileges, and no containers")
|
|
def test_loop(self):
|
|
"""Test the loop_* functions (requires loop support)"""
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img:
|
|
mksparse(disk_img.name, 42 * 1024**2)
|
|
loop_dev = loop_attach(disk_img.name)
|
|
try:
|
|
self.assertTrue(loop_dev is not None)
|
|
self.assertEqual(loop_dev[5:], get_loop_name(disk_img.name))
|
|
finally:
|
|
loop_detach(loop_dev)
|
|
|
|
@unittest.skipUnless(os.geteuid() == 0 and not os.path.exists("/.in-container"), "requires root privileges, and no containers")
|
|
def test_loop_context(self):
|
|
"""Test the LoopDev context manager (requires loop)"""
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img:
|
|
mksparse(disk_img.name, 42 * 1024**2)
|
|
with LoopDev(disk_img.name) as loop_dev:
|
|
self.assertTrue(loop_dev is not None)
|
|
self.assertEqual(loop_dev[5:], get_loop_name(disk_img.name))
|
|
|
|
@unittest.skipUnless(os.geteuid() == 0 and not os.path.exists("/.in-container"), "requires root privileges, and no containers")
|
|
def test_dm(self):
|
|
"""Test the dm_* functions (requires device-mapper support)"""
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img:
|
|
mksparse(disk_img.name, 42 * 1024**2)
|
|
with LoopDev(disk_img.name) as loop_dev:
|
|
self.assertTrue(loop_dev is not None)
|
|
dm_name = dm_attach(loop_dev, 42 * 1024**2)
|
|
try:
|
|
self.assertTrue(dm_name is not None)
|
|
finally:
|
|
dm_detach(dm_name)
|
|
|
|
@unittest.skipUnless(os.geteuid() == 0 and not os.path.exists("/.in-container"), "requires root privileges, and no containers")
|
|
def test_dmdev(self):
|
|
"""Test the DMDev context manager (requires device-mapper support)"""
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img:
|
|
mksparse(disk_img.name, 42 * 1024**2)
|
|
with LoopDev(disk_img.name) as loop_dev:
|
|
self.assertTrue(loop_dev is not None)
|
|
with DMDev(loop_dev, 42 * 1024**2) as dm_name:
|
|
self.assertTrue(dm_name is not None)
|
|
|
|
@unittest.skipUnless(os.geteuid() == 0 and not os.path.exists("/.in-container"), "requires root privileges, and no containers")
|
|
def test_mount(self):
|
|
"""Test the Mount context manager (requires loop)"""
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img:
|
|
mksparse(disk_img.name, 42 * 1024**2)
|
|
runcmd(["mkfs.ext4", "-L", "Anaconda", "-b", "4096", "-m", "0", disk_img.name])
|
|
with LoopDev(disk_img.name) as loopdev:
|
|
self.assertTrue(loopdev is not None)
|
|
with Mount(loopdev) as mnt:
|
|
self.assertTrue(mnt is not None)
|
|
|
|
@unittest.skipUnless(os.geteuid() == 0 and not os.path.exists("/.in-container"), "requires root privileges, and no containers")
|
|
def test_mkdosimg(self):
|
|
"""Test mkdosimg function (requires loop)"""
|
|
with tempfile.TemporaryDirectory(prefix="lorax.test.") as work_dir:
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img:
|
|
mkfakerootdir(work_dir)
|
|
mkdosimg(work_dir, disk_img.name)
|
|
self.assertTrue(os.path.exists(disk_img.name))
|
|
file_details = get_file_magic(disk_img.name)
|
|
self.assertTrue("FAT " in file_details, file_details)
|
|
|
|
@unittest.skipUnless(os.geteuid() == 0 and not os.path.exists("/.in-container"), "requires root privileges, and no containers")
|
|
def test_mkext4img(self):
|
|
"""Test mkext4img function (requires loop)"""
|
|
with tempfile.TemporaryDirectory(prefix="lorax.test.") as work_dir:
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img:
|
|
mkfakerootdir(work_dir)
|
|
graft = {work_dir+"/ext4test/": "./tests/pylorax/templates/install-cmd.tmpl"}
|
|
mkext4img(work_dir, disk_img.name, graft=graft)
|
|
self.assertTrue(os.path.exists(disk_img.name))
|
|
file_details = get_file_magic(disk_img.name)
|
|
self.assertTrue("ext2 filesystem" in file_details, file_details)
|
|
|
|
@unittest.skipUnless(os.geteuid() == 0 and not os.path.exists("/.in-container"), "requires root privileges, and no containers")
|
|
def test_small_mkext4img(self):
|
|
"""Test mkext4img error handling"""
|
|
with tempfile.TemporaryDirectory(prefix="lorax.test.") as work_dir:
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img:
|
|
mkfakerootdir(work_dir)
|
|
# Add a 8MiB file
|
|
with open(joinpaths(work_dir, "large-file"), "w") as f:
|
|
for _ in range(5):
|
|
f.write("A" * 1024**2)
|
|
graft = {work_dir+"/ext4test/": "./tests/pylorax/templates/install-cmd.tmpl"}
|
|
try:
|
|
mkext4img(work_dir, disk_img.name, graft=graft, size=5*1024**2)
|
|
except CalledProcessError as e:
|
|
self.assertTrue(e.stdout and "No space left on device" in e.stdout)
|
|
|
|
@unittest.skipUnless(os.geteuid() == 0 and not os.path.exists("/.in-container"), "requires root privileges, and no containers")
|
|
def test_mkbtrfsimg(self):
|
|
"""Test mkbtrfsimg function (requires loop)"""
|
|
with tempfile.TemporaryDirectory(prefix="lorax.test.") as work_dir:
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img:
|
|
mkfakerootdir(work_dir)
|
|
mkbtrfsimg(work_dir, disk_img.name)
|
|
self.assertTrue(os.path.exists(disk_img.name))
|
|
file_details = get_file_magic(disk_img.name)
|
|
self.assertTrue("BTRFS Filesystem" in file_details, file_details)
|
|
|
|
@unittest.skipUnless(os.geteuid() == 0 and not os.path.exists("/.in-container"), "requires root privileges, and no containers")
|
|
def test_mkhfsimg(self):
|
|
"""Test mkhfsimg function (requires loop)"""
|
|
with tempfile.TemporaryDirectory(prefix="lorax.test.") as work_dir:
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img:
|
|
mkfakerootdir(work_dir)
|
|
mkhfsimg(work_dir, disk_img.name, label="test")
|
|
self.assertTrue(os.path.exists(disk_img.name))
|
|
file_details = get_file_magic(disk_img.name)
|
|
self.assertTrue("Macintosh HFS" in file_details, file_details)
|
|
|
|
def test_default_image_name(self):
|
|
"""Test default_image_name function"""
|
|
for compression, suffix in [("xz", ".xz"), ("gzip", ".gz"), ("bzip2", ".bz2"), ("lzma", ".lzma")]:
|
|
filename = default_image_name(compression, "foobar")
|
|
self.assertTrue(filename.endswith(suffix))
|
|
|
|
@unittest.skipUnless(os.geteuid() == 0 and not os.path.exists("/.in-container"), "requires root privileges, and no containers")
|
|
def test_partition_mount(self):
|
|
"""Test PartitionMount context manager (requires loop)"""
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img:
|
|
self.assertTrue(mkfakediskimg(disk_img.name))
|
|
# Make sure it can mount the / with /etc/passwd
|
|
with PartitionMount(disk_img.name) as img_mount:
|
|
self.assertTrue(img_mount is not None)
|
|
self.assertTrue(os.path.isdir(img_mount.mount_dir))
|
|
self.assertTrue(os.path.exists(joinpaths(img_mount.mount_dir, "/etc/passwd")))
|
|
|
|
# Make sure submount works
|
|
with PartitionMount(disk_img.name, submount="/a-sub-mount/") as img_mount:
|
|
self.assertTrue(img_mount is not None)
|
|
self.assertTrue(os.path.isdir(img_mount.mount_dir))
|
|
self.assertTrue(os.path.exists(joinpaths(img_mount.mount_dir, "/etc/passwd")))
|
|
|
|
# Make sure it can mount the /boot partition with a custom mount_ok function
|
|
def mount_ok(mount_dir):
|
|
kernels = glob.glob(joinpaths(mount_dir, "vmlinuz-*"))
|
|
return len(kernels) > 0
|
|
|
|
with PartitionMount(disk_img.name, mount_ok=mount_ok) as img_mount:
|
|
self.assertTrue(img_mount is not None)
|
|
self.assertTrue(os.path.isdir(img_mount.mount_dir))
|
|
self.assertFalse(os.path.exists(joinpaths(img_mount.mount_dir, "/etc/passwd")))
|
|
self.assertTrue(os.path.exists(joinpaths(img_mount.mount_dir, "vmlinuz-4.18.13-200.fc28.x86_64")))
|
|
self.assertTrue(os.path.exists(joinpaths(img_mount.mount_dir, "initramfs-4.18.13-200.fc28.x86_64.img")))
|
|
|
|
@unittest.skipUnless(os.geteuid() == 0 and not os.path.exists("/.in-container"), "requires root privileges, and no containers")
|
|
def test_mkfsimage_from_disk(self):
|
|
"""Test creating a fsimage from the / partition of a disk image"""
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as disk_img:
|
|
self.assertTrue(mkfakediskimg(disk_img.name))
|
|
with tempfile.NamedTemporaryFile(prefix="lorax.test.disk.") as fs_img:
|
|
mkfsimage_from_disk(disk_img.name, fs_img.name)
|
|
self.assertTrue(os.path.exists(fs_img.name))
|
|
file_details = get_file_magic(fs_img.name)
|
|
self.assertTrue("ext2 filesystem" in file_details, file_details)
|