tools/run_guest.py: Support restricting CPU list available to guest

Signed-off-by: Hector Martin <marcan@marcan.st>

proxyclient/m1n1/utils.py

@@ -121,6 +121,17 @@ def unhex(s):
     s = re.sub(r"/\*.*?\*/", "", s)
     return bytes.fromhex(s.replace(" ", "").replace("\n", ""))
 
+def parse_indexlist(s):
+    items = set()
+    for i in s.split(","):
+        if "-" in i:
+            a, b = map(int, i.split("-", 1))
+            for i in range(a, b + 1):
+                items.add(i)
+        else:
+            items.add(int(i))
+    return items
+
 FourCC = ExprAdapter(Int32ul,
                      lambda d, ctx: d.to_bytes(4, "big").decode("latin-1"),
                      lambda d, ctx: int.from_bytes(d.encode("latin-1"), "big"))

proxyclient/tools/run_guest.py

@@ -13,6 +13,7 @@ parser.add_argument('-S', '--shell', action="store_true")
 parser.add_argument('-e', '--hook-exceptions', action="store_true")
 parser.add_argument('-d', '--debug-xnu', action="store_true")
 parser.add_argument('-l', '--logfile', type=pathlib.Path)
+parser.add_argument('-C', '--cpus', default=None)
 parser.add_argument('payload', type=pathlib.Path)
 parser.add_argument('boot_args', default=[], nargs="*")
 args = parser.parse_args()
@@ -34,6 +35,18 @@ hv.hook_exceptions = args.hook_exceptions
 
 hv.init()
 
+if args.cpus:
+    avail = [i.name for i in hv.adt["/cpus"]]
+    want = set(f"cpu{i}" for i in args.cpus)
+    for cpu in avail:
+        if cpu in want:
+            continue
+        try:
+            del hv.adt[f"/cpus/{cpu}"]
+            print(f"Disabled {cpu}")
+        except KeyError:
+            continue
+
 if args.debug_xnu:
     hv.adt["chosen"].debug_enabled = 1