build_disk_image.py

#!/usr/bin/env python3
#
# Builds a disk image from individual partition images and a partition
# table description. The actual disk image is wrapped up into a tar file
# because the raw image file is sparse.
#
# The input partition images are also expected to be given as tzst files,
#
# Call example:
# build_disk_image -p partitions.csv -o disk.img.tar part1.tzst part2.tzst ...
#
importargparse
importos
importsubprocess
importsys
importtarfile


defread_partition_description(data):
gpt_entries= []
forlineindata.split("\n"):
ifline.startswith("#") ornotline:
continue
cols=tuple(map(lambdas: s.strip(), line.split(",")))
gpt_entries.append(
 {
"name": cols[0],
"start": int(cols[1]),
"size": int(cols[2]),
"type": cols[3],
"uuid": cols[4],
"description": cols[5],
 }
 )
returngpt_entries


defvalidate_partition_table(gpt_entries):
"""
 Validate partition table.

 Validate that partitions do not overlap, start at suitable
 boundaries, and have sane types.
 """
end=2048
forentryingpt_entries:
ifentry["start"] <end:
raiseRuntimeError("Partition %s overlaps with previous"%entry["name"])
if (entry["start"] %2048) !=0:
raiseRuntimeError("Partition %s start is not aligned to 1MB boundary"%entry["name"])
if (entry["size"] %2048) !=0:
raiseRuntimeError("Partition %s size is not aligned to 1MB boundary"%entry["name"])
ifnotentry["type"] in ("U", "L", "M"):
raiseRuntimeError("Partition %s has unsupported type"%entry["name"])
end=entry["start"] +entry["size"]


defgenerate_sfdisk_script(gpt_entries):
type_map= {
# UEFI partition type
"U": "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
# Linux partition type
"L": "0FC63DAF-8483-4772-8E79-3D69D8477DE4",
# Microsoft basic data
"M": "EBD0A0A2-B9E5-4433-87C0-68B6B72699C7",
 }
lines= ["label: gpt", "label-id: 2B110BB7-CDEC-7D41-B97E-893EDCBE5428"]
forentryingpt_entries:
lines.append(
"start=%d,size=%d,type=%s,uuid=%s"% (entry["start"], entry["size"], type_map[entry["type"]], entry["uuid"])
 )
return"\n".join(lines) +"\n"


defprepare_diskimage(gpt_entries, image_file):
last=gpt_entries[-1]
image_size= (last["start"] +last["size"] +2048) *512

os.close(os.open(image_file, os.O_CREAT|os.O_RDWR|os.O_CLOEXEC|os.O_EXCL, 0o600))
os.truncate(image_file, image_size)

withsubprocess.Popen(["/usr/sbin/sfdisk", image_file], stdin=subprocess.PIPE) asproc:
proc.stdin.write(generate_sfdisk_script(gpt_entries).encode("utf-8"))
proc.stdin.close()
proc.wait()
ifproc.returncode!=0:
raiseRuntimeError("Build of partition table failed")


def_copyfile(source, target, size):
whilesize:
count=min(16*1024, size)
data=source.read(count)
target.write(data)
size-=len(data)


defwrite_partition_image_from_tzst(gpt_entry, image_file, partition_tzst):
tmpdir=os.getenv("ICOS_TMPDIR")
ifnottmpdir:
raiseRuntimeError("ICOS_TMPDIR env variable not available, should be set in BUILD script.")

partition_tf=os.path.join(tmpdir, "partition.tar")
subprocess.run(["zstd", "-q", "--threads=0", "-f", "-d", partition_tzst, "-o", partition_tf], check=True)

partition_tf=tarfile.open(partition_tf, mode="r:")
base=gpt_entry["start"] *512
withos.fdopen(os.open(image_file, os.O_RDWR), "wb+") astarget:
formemberinpartition_tf:
ifmember.path!="partition.img":
continue
ifmember.size>gpt_entry["size"] *512:
raiseRuntimeError("Image too large for partition %s"%gpt_entry["name"])
source=partition_tf.extractfile(member)
ifmember.type==tarfile.GNUTYPE_SPARSE:
foroffset, sizeinmember.sparse:
ifsize==0:
continue
source.seek(offset)
target.seek(offset+base)
_copyfile(source, target, size)
else:
target.seek(base)
_copyfile(source, target, member.size)


defselect_partition_file(name, partition_files):
forpartition_fileinpartition_files:
ifnameinos.path.basename(partition_file):
returnpartition_file

returnNone


defmain():
parser=argparse.ArgumentParser()
parser.add_argument("-o", "--out", help="Target (tar) file to write disk image to", type=str)
parser.add_argument("-p", "--partition_table", help="CSV file describing the partition table", type=str)
parser.add_argument("-s", "--expanded-size", help="Optional size to grow the image to", required=False, type=str)
parser.add_argument(
"partitions",
metavar="partition",
type=str,
nargs="*",
help="Partitions to write. These must match the CSV partition table entries.",
 )
parser.add_argument("--dflate", help="Path to our dflate tool", type=str)

args=parser.parse_args(sys.argv[1:])

out_file=args.out
partition_desc_file=args.partition_table
partition_files=list(args.partitions)

withopen(partition_desc_file, "r") asf:
gpt_entries=read_partition_description(f.read())
validate_partition_table(gpt_entries)

tmpdir=os.getenv("ICOS_TMPDIR")
ifnottmpdir:
raiseRuntimeError("ICOS_TMPDIR env variable not available, should be set in BUILD script.")

ifargs.dflate:
disk_image=os.path.join(tmpdir, "disk.img")
else:
# Disk optimization. If no dflate program is specified, we can
# simply attack the target file (out_file) directly, saving gigabytes
# of writes to disk.
disk_image=out_file
prepare_diskimage(gpt_entries, disk_image)

forentryingpt_entries:
# Skip over any partitions starting with "B_". These are empty in our
# published images, and stay this way until a live system upgrades
# into them.
ifentry["name"].startswith("B_"):
continue

# Remove the "A_" prefix from any partitions before doing a lookup.
prefix="A_"
name=entry["name"]
ifname.startswith(prefix):
name=name[len(prefix) :]

partition_file=select_partition_file(name, partition_files)

ifpartition_file:
write_partition_image_from_tzst(entry, disk_image, partition_file)
else:
print("No partition file for '%s' found, leaving empty"%name)

# Provide additional space for vda10, the final partition, for immediate QEMU use
ifargs.expanded_size:
subprocess.run(["truncate", "--size", args.expanded_size, disk_image], check=True)

# We use our tool, dflate, to quickly create a sparse, deterministic, tar.
# If dflate is ever misbehaving, it can be replaced with:
# tar cf <output> --sort=name --owner=root:0 --group=root:0 --mtime="UTC 1970-01-01 00:00:00" --sparse --hole-detection=raw -C <context_path> <item>
ifargs.dflate:
subprocess.run(
 [
args.dflate,
"--input",
disk_image,
"--output",
out_file,
 ],
check=True,
 )


if__name__=="__main__":
main()