168 lines
3.3 KiB
Python
Executable File
168 lines
3.3 KiB
Python
Executable File
#!/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.write('\n')
|
|
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 _GenerateSSHKey(self, root):
|
|
self._Exec(
|
|
'ssh-keygen',
|
|
'-f', os.path.join(root, 'ssh_host_ed25519_key'),
|
|
'-N', '',
|
|
'-t', 'ed25519',
|
|
)
|
|
|
|
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)
|
|
self._GenerateSSHKey(root)
|
|
|
|
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()
|