SystemID image building script

systemid/image.py

@@ -0,0 +1,157 @@
+#!/usr/bin/python3
+
+import argparse
+import json
+import os
+import random
+import shutil
+import subprocess
+import sys
+import tempfile
+import time
+
+
+parser = argparse.ArgumentParser(description='iconograph image')
+parser.add_argument(
+    '--db-dir',
+    dest='db_dir',
+    action='store',
+    required=True)
+parser.add_argument(
+    '--device',
+    dest='device',
+    action='store',
+    required=True)
+FLAGS = parser.parse_args()
+
+
+class Imager(object):
+
+  _ALPHABET = '0123456789ABCDEF'
+  _LEN = 6
+
+  def __init__(self, device, db_dir):
+    self._device = device
+    self._db_dir = db_dir
+
+  def _Exec(self, *args, **kwargs):
+    print('+', args)
+    subprocess.check_call(args, **kwargs)
+
+  def _PartDev(self, part_num):
+    time.sleep(1)
+    args = {
+      'device': self._device,
+      'part': part_num,
+    }
+    options = [
+      '%(device)s%(part)d' % args,
+      '%(device)sp%(part)d' % args,
+    ]
+    while True:
+      for option in options:
+        if os.path.exists(option):
+          return option
+
+  def _PartitionAndMkFS(self):
+    self._Exec(
+        'parted',
+        '--script',
+        self._device,
+        'mklabel', 'msdos')
+
+    self._Exec(
+        'parted',
+        '--script',
+        '--align', 'optimal',
+        self._device,
+        'mkpart', 'primary', 'ext4', '0%', '100%')
+    self._Exec(
+        'mkfs.ext4',
+        '-L', 'SYSTEMID',
+        '-F',
+        self._PartDev(1))
+
+  def _Mount(self):
+    root = tempfile.mkdtemp()
+    self._rmtree.append(root)
+
+    self._Exec(
+        'mount',
+        self._PartDev(1),
+        root)
+    self._umount.append(root)
+
+    return root
+
+  def _ChooseID(self):
+    ret = []
+    for _ in range(self._LEN):
+      ret.append(random.choice(self._ALPHABET))
+    return ''.join(ret)
+
+  def _SaveID(self, new_id):
+    db_path = os.path.join(self._db_dir, 'systemid.json')
+    try:
+      with open(db_path, 'r') as fh:
+        db = json.load(fh)
+    except FileNotFoundError:
+      db = {
+        'ids': [],
+      }
+    if new_id in db['ids']:
+      return False
+    db['ids'].append(new_id)
+    with tempfile.NamedTemporaryFile(dir=self._db_dir, delete=False, mode='w') as fh:
+      try:
+        json.dump(db, fh)
+        fh.flush()
+        os.rename(fh.name, db_path)
+      except Exception:
+        os.unlink(fh.name)
+        raise
+    return True
+
+  def _WriteID(self, root, new_id):
+    with open(os.path.join(root, 'systemid'), 'w') as fh:
+      fh.write("""# Autogenerated by iconograph/systemid/image.py
+SYSTEMID=%(system_id)s
+""" % {
+        'system_id': new_id,
+      })
+
+  def _Image(self):
+    self._PartitionAndMkFS()
+    root = self._Mount()
+    while True:
+      new_id = self._ChooseID()
+      if self._SaveID(new_id):
+        break
+    self._WriteID(root, new_id)
+    print("""
+
+==============
+New ID: \033[91m%s\033[00m
+==============
+
+""" % new_id)
+
+  def Image(self):
+    self._umount = []
+    self._rmtree = []
+    try:
+      self._Image()
+    finally:
+      for path in reversed(self._umount):
+        self._Exec('umount', path)
+      for path in self._rmtree:
+        shutil.rmtree(path)
+
+
+def main():
+  imager = Imager(FLAGS.device, FLAGS.db_dir)
+  imager.Image()
+
+
+if __name__ == '__main__':
+  main()