#! @python@/bin/python -B
import base64
import json
import os
import shutil
import socket
import subprocess
import sys

import yaml

CONF_FILE = '/etc/qclk/config.yaml'
WG_KEY_FILE = '/etc/qclk/wg.key'

IWD_PATH = '/var/lib/iwd'
IWD_CONF = '''[Settings]
Hidden={hidden}

[Security]
Passphrase={psk}
'''

MANAGEMENT_NET = '''[Match]
Name=management

[Network]
Address={ip}/32
'''

def log(m):
  print(m, file=sys.stderr)

def iwd_ssid_basename(ssid: str) -> str:
  if ssid.isalnum() and not any(c in ssid for c in '-_ '):
    return ssid

  return f"={ssid.encode('utf-8').hex()}"

class Configurer:
  tmpdir = '/tmp/qclk'

  def __init__(self, conf_file: str, wg_key_file: str):
    with open(conf_file) as f:
      self.config = yaml.safe_load(f)

    with open(wg_key_file, 'rb') as f:
      output = subprocess.check_output(['@wireguardTools@/bin/wg', 'pubkey'], input=f.read())
    pubkey = base64.b64decode(output)
    self.id = pubkey[:4].hex()

    os.makedirs(self.tmpdir, exist_ok=True)

  def _setup_hostname(self):
    hostname = f'qclk-{self.id}'
    log(f"Setting hostname to '{hostname}'")
    socket.sethostname(hostname)

  def _setup_wifi(self):
    os.makedirs(IWD_CONF, exist_ok=True)
    if 'wifi' in self.config:
      conf = self.config['wifi']
      tmp = os.path.join(self.tmpdir, 'wifi.psk')
      with open(tmp, 'w') as f:
        print(
          IWD_CONF.format(
            hidden=str(conf['hidden']).lower(),
            psk=conf['password']),
          file=f)

      fname = f"{iwd_ssid_basename(conf['ssid'])}.psk"
      log(f"Writing IWD config file '{fname}' for network '{conf['ssid']}'")
      shutil.move(tmp, os.path.join(IWD_PATH, fname))
    else:
      fname = None

    for f in os.listdir(IWD_PATH):
      if (fname is not None and f != fname) and f.endswith('.psk'):
        log(f"Cleaning up old IWD config '{f}'")
        os.remove(os.path.join(IWD_PATH, f))

  def _setup_mgmt(self):
    conf = self.config['management']
    tmp = os.path.join(self.tmpdir, 'management.network')
    with open(tmp, 'w') as f:
      print(MANAGEMENT_NET.format(ip=conf['ip']), file=f)

    log(f"Configuring management IP {conf['ip']}")
    os.makedirs('/run/systemd/network', exist_ok=True)
    shutil.move(tmp, '/run/systemd/network/20-management.network')
    subprocess.check_call(['@systemd@/bin/networkctl', 'reload'])

  def reconcile(self):
    self._setup_hostname()
    self._setup_wifi()
    self._setup_mgmt()

def main():
  c = Configurer(CONF_FILE, WG_KEY_FILE)
  c.reconcile()

if __name__ == '__main__':
  main()