From da5150084e87cf70c7fdf7f8869578b903707432 Mon Sep 17 00:00:00 2001
From: Christian Kellner <christian@kellner.me>
Date: Tue, 12 Oct 2021 20:53:37 +0000
Subject: [PATCH] stages/parted: new stage to partition a device

Like the existing sfdisk stage, the parted stage can be used to create
a partition table on the specified device. In contrast to sfdisk, it
does not support uuids, neither for specifying the partition types,
nor the actual uuid of the partition. The current implementation only
supports GPT.
This stage is meant to be used on older systems, like RHEL 7, where
sfdisk exists but does not support GPT (or --json).
---
 stages/org.osbuild.parted | 165 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 165 insertions(+)
 create mode 100755 stages/org.osbuild.parted

diff --git a/stages/org.osbuild.parted b/stages/org.osbuild.parted
new file mode 100755
index 00000000..b01997bc
--- /dev/null
+++ b/stages/org.osbuild.parted
@@ -0,0 +1,165 @@
+#!/usr/bin/python3
+"""
+Partition a target using parted(8)
+
+Use the `parted`(8) command line utility to partition the specified
+`device`.
+
+Buildhost commands used: `parted`.
+"""
+
+import subprocess
+import sys
+
+
+import osbuild.api
+
+
+SCHEMA_2 = r"""
+"devices": {
+  "type": "object",
+  "additionalProperties": true,
+  "required": ["device"],
+  "properties": {
+    "device": {
+      "type": "object",
+      "additionalProperties": true
+    }
+  }
+},
+"options": {
+  "additionalProperties": false,
+  "required": ["label"],
+  "properties": {
+    "label": {
+      "description": "The type of the partition table",
+      "type": "string",
+      "enum": ["gpt"]
+    },
+    "partitions": {
+      "description": "Partition layout ",
+      "type": "array",
+      "required": ["start", "size", "name"],
+      "items": {
+        "description": "Description of one partition",
+        "type": "object",
+        "properties": {
+          "bootable": {
+            "description": "Mark the partition as bootable",
+            "type": "boolean"
+          },
+          "name": {
+            "description": "The partition name (GPT)",
+            "type": "string"
+          },
+          "size": {
+            "description": "The size of this partition",
+            "type": "integer"
+          },
+          "start": {
+            "description": "The start offset of this partition",
+            "type": "integer"
+          },
+          "type": {
+            "description": "The partition type",
+            "type": "string"
+          }
+        }
+      }
+    }
+  }
+}
+"""
+
+
+class Partition:
+    def __init__(self,
+                 pttype: str = None,
+                 start: int = None,
+                 size: int = None,
+                 bootable: bool = False,
+                 name: str = None,
+                 uuid: str = None):
+        self.type = pttype
+        self.start = start
+        self.size = size
+        self.bootable = bootable
+        self.name = name
+        self.uuid = uuid
+        self.index = None
+
+    @property
+    def start_in_bytes(self):
+        return (self.start or 0) * 512
+
+    @property
+    def size_in_bytes(self):
+        return (self.size or 0) * 512
+
+    @classmethod
+    def from_json(cls, js) -> "Partition":
+        p = cls(pttype=js.get("type"),
+                start=js.get("start"),
+                size=js.get("size"),
+                bootable=js.get("bootable"),
+                name=js.get("name"),
+                uuid=js.get("uuid"))
+        return p
+
+
+class PartitionTable:
+    def __init__(self, label, uuid, partitions):
+        self.label = label
+        self.uuid = uuid
+        self.partitions = partitions or []
+
+    def __getitem__(self, key) -> Partition:
+        return self.partitions[key]
+
+    def write_to(self, target):
+        """Write the partition table to disk"""
+        # generate the command for parted to create the table
+        commands = [
+            "unit", "s",
+            "mklabel", self.label
+        ]
+
+        for i, p in enumerate(self.partitions):
+
+            commands += [
+                "mkpart", f'"{p.name}"', str(p.start), str(p.start + p.size - 1)
+            ]
+
+            if p.bootable:
+                commands += ["set", str(i+1), "boot", "on"]
+
+            if p.type:
+                commands += ["set", str(i+1), p.type, "on"]
+
+        subprocess.run(["parted", "-a", "none", "-s",
+                        target, "--"] + commands,
+                       encoding='utf-8',
+                       check=True)
+
+        subprocess.run(["parted", "-a", "none", "-s",
+                        target, "--", "print"],
+                       encoding='utf-8',
+                       check=True)
+
+
+def main(devices, options):
+    device = devices["device"]["path"]
+
+    pttype = options["label"]
+    partitions = options.get("partitions")
+
+    parts = [Partition.from_json(p) for p in partitions]
+    pt = PartitionTable(pttype, None, parts)
+
+    pt.write_to(device)
+
+
+if __name__ == '__main__':
+    args = osbuild.api.arguments()
+    ret = main(args["devices"], args["options"])
+    sys.exit(ret)