#!/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 _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()