import sys
import tempfile
import shutil
import os
import os.path

# For some weird reason cloud-init won't work with files dumped by pyyaml, and I'm fucking tired of debugging this
cloudinit = """#cloud-config
hostname: {hostname}
fqdn: {hostname}.local
manage_etc_hosts: true
users:
  - name: ubuntu
    sudo: ALL=(ALL) NOPASSWD:ALL
    groups: users, admin, sudo
    home: /home/ubuntu
    shell: /bin/bash
    lock_passwd: false
    ssh-authorized-keys:
      - ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAACAQDvNCyaO5P/vIuzxuiBY9XJ9YrZXByLUOUGpEjIdw8G1s2f2Iu8ung7S6x6XnCckaclTk3zJgg4cJlSyWuEUeMrz4npKvd3W8P4n8lj8xwhqDQ9ywta9TXriBih5Z/PPNvAYncx28F9BIDA71Q0Bybmv4/XTtPPgxE6hpjmVnjVE5pgwKceVR3Cj4Yl6+at3bC3Z15K61sNmciCV7yGUOpDkf3ZS/VBDOQIZ1GVEvkMBhNiYw2V/X2EwY3EBHVFjiryeMiwCcutTfsDvCXxsMr/1KkOu+fAbcS77+T8n7igaTQ1P59hH6VsUhwemIpnwt5NzPa9fl/+Nedw7WXx0GB1OZDUorH3L5cdKQ+nDEjUAuL+F5jjHHJMYk6V4HqNzuBxZ6lbUzdin0jyUoW8NDxDBFp+J18Mid4QiQmPhLJPkRCjY/Mf0uRnnjIB0HtysBRkHgnwVdYUu52eY/QFHqiAra2kvBspp4MiryKGpZbKTF81DnibDOrWWDB12x6+HXnY6Gpbyk0JeeILgh2VD1L0CRnyg4VC9sSUEzOnHi+4X99D89DUuvthlZWcFTny83IhX44dyI82YTNzlET6oJss76mHNLBG+6k80bVHvPw4qLNAIlpX905EaJ3SQhMonw7geOasLzT4IgSLbA22jYzKcYU8zslCBykr50jDJd3u7Q==
      - ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAACAQCvzoxDG/eJQPNOUWkdzRzSk9oq8kxrB5x9UdlGQiARlBIMNiTZhvR7KxvYzaGxIQlgk+/4rikIcoQyPMz0yiQA8spvZv7QyG+ikhrpHxly2iHr9lFI4nFqCXfMUH+QrRElB14f/Wper80gtUnuEkH8azz9MdD7RrJ8jvyzIXsEN71enH58xw5jsGFiyEHaNmSJyuQwK/gcRAborne4gcFHeUV0AZB0imLB9c3Jlif8yB/4n6GUu+am/70VOKzv0aeM0vgMvZOCYyyYWV7E4n/A0tzn+BV/ioNmhA1HrH5s+yNc+VNMPyyH69FG2r4BQaefqhtLNJNb+W1MAJPgLW9DtqSQ1f5vSkw1yNY6wPZAkb/Y+Yy+APtwXaszcWTPFfFiZuzde1NLJIbi6pAGtRmpn0FJFiBcqIgAiJciuCLwYxkNASzyqVcHiFfaYK1MbGYv8vLGhw5ustA3EZpzNDkeHUr2zvWcL1frM3Ku4p08vRyeAXJqteZ7zn0r9mi5hu3LXx/F219UYaMu5Hy/Gcst2PedYeyjGfXhhaFTKpdTjquCWvY0+6TL4X0Stwn737OC//xQOlzL/Db5nU/q5OJfVyAfh/uKuLjB383vRPibegDrybbsThfPnXFPL9NbDIOOHvGwlHpGDG66KTDKI++jaeO9s6IsW9Ng72x7jtlNWw==
      - ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAACAQD3bnVGHfIyBbceOXuD4QgYET/lZyXYQY89vewBCmR97Z7/hgCEydjPM7HewLsJaRqlueu9Tc0dlvXAeen2BU2WgRrZV0B4WaFzLCf1FVpeWz7t6EaJl37Zuuy5ActJS2ECqe7El5Od15sS+DxW3AJ24lcy8eFsC4cbQt8v5S/76gE2D7NUEGTxDGTk154NRr8Kr/P5Fua9gjx5z8UUXbcKBAhbYGtp+Iv+htwAvyRptPgpAtQt/j7nQUSVNZ+d7j6OJjtquh8YDxf5faYzusLF3JeDOr64QUrvswpwBB4Qj5VrYqZgQBEDERB69KexeTGHeFLM+xu5Ls0Oo2+owiAZf8BPEfVMLzner4VhVuE+VGc4+NKf2zvx/pe8qRgWgliKSbCQFhenIX7NbORY/J8376p1hpANGB9Q0Lt2JT/Wx0qha9EMYQV+gxMWWdhgxpxGrze8flzd1KH1GtVIGRotIv65ecjsXwQcpp/ke7DOog5XcaGaasWJ+5oWUvWp7riI9j95nn1Gs+D2ccF45oClEHv5zP9XxhlHicqt1OLuRbmYWp4u6lndCnvvVtaAfPfnU+tdRrze6QqInPHEsWEN5kVUR7WfZXozMdi6XzZiD/rSZ9kocCRg/LmNfy7xLqIwYqONU+snbty9x2wkaA8YQEsPdWM3PmWKztBEOPiGKw==
      - ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAACAQCno5mh0+xtkCXoeSt5CrXv8IR1/q+FgiZQeZ4rIqKB1Jko3pY0TiYaHv8hZZ3g/aFfuM8expMrhT5BDvSeEYZ15JkBzBcxo/Hyup6IloOYaw1p1v1GdNo6M0V5cC/B7dyGHKcVllzDRLN6kKsWKAyTgtSd8v06sz9FLJe4I9jOkqKcOTjgUfAvF98fCVIo0eOeWFfA/dDXaW2Tu9xaw3sHLnJFDd+kp1r99d7/Tgmf1hITQaVQ/+L6zNxpR4zWzfdLPvKphfVB34VE6Glkn+PnTGV7Q32+VbgcPizbGVcR7V2EHc53u0evB5BkLwfmDQZRKwd5pS5QHFEnPVwIecaJsYCNDFespj2BwJWiiWvL2cOlU15g0r4uj9+nbzsBPdYNJiXD3VZVEGMvTUb9SkmnMlQl10x0a7j00qaPhJ2WROtYemioC9y7MLnHheotjfdaaAmf46EN5N1c08ULbaqOt1AENLdM+PAUeVfqzI2hfQVuTFM+hPdKFWo8ik/SlztbCjL8a6Isrr4ZzWrJVoqg32PnucSvmIRtCGrj5Lt5WViYTAy9SjrQwQ0jdeOtUYyzkeIwpJSmbRLHILL5zLAJFtIKND+XgRclzIM7uD2UqtzEGtmuLE8qx05bukwa7nqjG6UwDpN3nVO3aXcjnzLyuHWnz2NJc0ye2oVwdqPzZQ==
      - ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAACAQCfNP0DqG5ULSZVC6YUsS5NPfkSVwIOwuIoPJPAaQOfj6/+ZXXgb59x/X9gzfQ5peMhpa1m7z0qowCb+2gvDa07/2RwDEjTZzHWPYCxHaQPIRG8ttTY3B0PM5ftJlvjLtuPyeGeXU1L6f9s8pmWg855371HnyMmFzO2iA1G2FxIIZmWOICh9fZ/LeUwpmuRNuCovi7lZ3v3QCgFm1uYMbbCBOd7aF9H6nigFf5PuwjA+JkrD1t2MGscHHt9GVs2KMaNIfqdZ7jiLstcyKe0E/ADD+fXGBSkhLFUHUL88oO4rr13fzGGNmZmiF+ESjAXWuqR4VjC9rCErAXIiBI07QGFzEpVMvSoZVIeqrbKBSo9T8pptsYYQ6DzZ5F4UanepnrEPiabyi4GtZrqc3dDBinGQ25UBiFFDtJ7cogrDppPopE/ExtjNZG/Lbyshqnltfj2dKgxSyMqVQFyfIR3UIKNscS1uo+2yKwpOE0LSWfAaa5U36UdYjx9LCL0ODdA1kwKWoqUEbGRjX7slfrKzrn5zc1faOYonZyI3dmNPcYXz6dqY3DyHhykszdO939z5+GtjiTpen7CuAl8QPm6Ad0wvONFzNcYFtN40UHmiUPkGxMNWsDJWuPWGpR1IDDT8oYT0WQXJyTTQiHP+0wVc8CQI8Xjw7ivEYLJaQeP+UNiYQ==

# only cert auth via ssh (console access can still login)
ssh_pwauth: false
disable_root: true
chpasswd:
  list: |
     ubuntu:ansibleme
  expire: False
packages:
  - qemu-guest-agent
  - python3
# written to /var/log/cloud-init-output.log
final_message: "The system is finally up, after $UPTIME seconds"
"""

libvirt = """<domain type="kvm">
  <name>{hostname}</name>
  <metadata>
    <libosinfo:libosinfo xmlns:libosinfo="http://libosinfo.org/xmlns/libvirt/domain/1.0">
      <libosinfo:os id="http://ubuntu.com/ubuntu/20.04"/>
    </libosinfo:libosinfo>
  </metadata>
  <memory unit="MiB">{ram}</memory>
  <currentMemory unit="MiB">{ram}</currentMemory>
  <vcpu placement="static">2</vcpu>
  <os>
    <type arch="x86_64" machine="pc-q35-6.1">hvm</type>
    <boot dev="hd"/>
  </os>
  <features>
    <acpi/>
    <apic/>
    <vmport state="off"/>
  </features>
  <cpu mode="host-model" check="partial"/>
  <clock offset="utc">
    <timer name="rtc" tickpolicy="catchup"/>
    <timer name="pit" tickpolicy="delay"/>
    <timer name="hpet" present="no"/>
  </clock>
  <on_poweroff>destroy</on_poweroff>
  <on_reboot>restart</on_reboot>
  <on_crash>destroy</on_crash>
  <pm>
    <suspend-to-mem enabled="no"/>
    <suspend-to-disk enabled="no"/>
  </pm>
  <devices>
    <emulator>/usr/bin/qemu-system-x86_64</emulator>
    <disk type="file" device="cdrom">
      <driver name="qemu" type="raw"/>
      <source file="{cloudinitiso}"/>
      <target dev="sda" bus="sata"/>
      <readonly/>
    </disk>
    <disk type="file" device="disk">
      <driver name="qemu" type="qcow2"/>
      <source file="{hostdrive}"/>
      <target dev="vda" bus="virtio"/>
    </disk>
    <controller type="usb" index="0" model="qemu-xhci" ports="15">
      <address type="pci" domain="0x0000" bus="0x02" slot="0x00" function="0x0"/>
    </controller>
    <controller type="scsi" index="0" model="virtio-scsi">
      <address type="pci" domain="0x0000" bus="0x03" slot="0x00" function="0x0"/>
    </controller>
    <controller type="sata" index="0">
      <address type="pci" domain="0x0000" bus="0x00" slot="0x1f" function="0x2"/>
    </controller>
    <controller type="pci" index="0" model="pcie-root"/>
    <controller type="pci" index="1" model="pcie-root-port">
      <model name="pcie-root-port"/>
      <target chassis="1" port="0x10"/>
      <address type="pci" domain="0x0000" bus="0x00" slot="0x02" function="0x0" multifunction="on"/>
    </controller>
    <controller type="pci" index="2" model="pcie-root-port">
      <model name="pcie-root-port"/>
      <target chassis="2" port="0x11"/>
      <address type="pci" domain="0x0000" bus="0x00" slot="0x02" function="0x1"/>
    </controller>
    <controller type="pci" index="3" model="pcie-root-port">
      <model name="pcie-root-port"/>
      <target chassis="3" port="0x12"/>
      <address type="pci" domain="0x0000" bus="0x00" slot="0x02" function="0x2"/>
    </controller>
    <controller type="pci" index="4" model="pcie-root-port">
      <model name="pcie-root-port"/>
      <target chassis="4" port="0x13"/>
      <address type="pci" domain="0x0000" bus="0x00" slot="0x02" function="0x3"/>
    </controller>
    <controller type="pci" index="5" model="pcie-root-port">
      <model name="pcie-root-port"/>
      <target chassis="5" port="0x14"/>
      <address type="pci" domain="0x0000" bus="0x00" slot="0x02" function="0x4"/>
    </controller>
    <controller type="pci" index="6" model="pcie-root-port">
      <model name="pcie-root-port"/>
      <target chassis="6" port="0x15"/>
      <address type="pci" domain="0x0000" bus="0x00" slot="0x02" function="0x5"/>
    </controller>
    <controller type="pci" index="7" model="pcie-root-port">
      <model name="pcie-root-port"/>
      <target chassis="7" port="0x16"/>
      <address type="pci" domain="0x0000" bus="0x00" slot="0x02" function="0x6"/>
    </controller>
    <controller type="pci" index="8" model="pcie-root-port">
      <model name="pcie-root-port"/>
      <target chassis="8" port="0x17"/>
      <address type="pci" domain="0x0000" bus="0x00" slot="0x02" function="0x7"/>
    </controller>
    <controller type="virtio-serial" index="0">
      <address type="pci" domain="0x0000" bus="0x04" slot="0x00" function="0x0"/>
    </controller>
{interfacecfg}
    <serial type="pty">
      <target type="isa-serial" port="0">
        <model name="isa-serial"/>
      </target>
    </serial>
    <console type="pty">
      <target type="serial" port="0"/>
    </console>
    <channel type="unix">
      <target type="virtio" name="org.qemu.guest_agent.0"/>
      <address type="virtio-serial" controller="0" bus="0" port="1"/>
    </channel>
    <channel type="spicevmc">
      <target type="virtio" name="com.redhat.spice.0"/>
      <address type="virtio-serial" controller="0" bus="0" port="2"/>
    </channel>
    <input type="tablet" bus="usb">
      <address type="usb" bus="0" port="1"/>
    </input>
    <input type="mouse" bus="ps2"/>
    <input type="keyboard" bus="ps2"/>
    <graphics type="spice" autoport="yes">
      <listen type="address"/>
    </graphics>
    <audio id="1" type="spice"/>
    <video>
      <model type="qxl" ram="65536" vram="65536" vgamem="16384" heads="1" primary="yes"/>
      <address type="pci" domain="0x0000" bus="0x00" slot="0x01" function="0x0"/>
    </video>
    <redirdev bus="usb" type="spicevmc">
      <address type="usb" bus="0" port="2"/>
    </redirdev>
    <redirdev bus="usb" type="spicevmc">
      <address type="usb" bus="0" port="3"/>
    </redirdev>
    <memballoon model="virtio">
      <address type="pci" domain="0x0000" bus="0x06" slot="0x00" function="0x0"/>
    </memballoon>
    <rng model="virtio">
      <backend model="random">/dev/urandom</backend>
      <address type="pci" domain="0x0000" bus="0x07" slot="0x00" function="0x0"/>
    </rng>
  </devices>
</domain>
"""


def run_iso(outdir:str, hostname:str, ipaddr:str):
  isoout = os.path.join(outdir, 'iso')
  
  fhandle, fname = tempfile.mkstemp(suffix='.yaml')
  with open(fhandle, "wt") as of:
	  of.write(
      cloudinit.format(hostname=hostname)
    )
  
  os.makedirs(isoout, exist_ok=True)
  os.system(f'cloud-localds -v {isoout}/cloudinit-{hostname}.iso {fname}')
  os.unlink(fname)

def run_img(outdir:str, hostname:str, img_src:str):
  imgout = os.path.join(outdir, 'img')
  os.makedirs(imgout, exist_ok=True)

  hostdrive = os.path.join(imgout, f"{hostname}.qcow2")

  shutil.copy(img_src, hostdrive)
  os.system(f"qemu-img resize {hostdrive} +100G")
  

def run_xml(outdir:str, hostname:str, ram:int, net:str, ipaddr:str):
  xmlout = os.path.join(outdir, 'xml')
  isoout = os.path.join(outdir, 'iso')
  cloudinitiso = os.path.join(isoout, f"cloudinit-{hostname}.iso")
  imgout = os.path.join(outdir, 'img')
  hostdrive = os.path.join(imgout, f"{hostname}.qcow2")
  os.makedirs(xmlout, exist_ok=True)

  interfacecfg = ""

  interfaces=list(net.split(','))

  for i, n in enumerate(interfaces):

    multi = ""
    if i == 0 and len(interfaces) > 1:
      multi = 'multifunction="on"'

    # Defining address like this fixes the interface to be enp1s0
    # And consequitive interfaces will be enp1s0f{n}
    interfacecfg += f"""    <interface type="network">
        <source network="{n}"/>
        <model type="virtio"/>
        <address type="pci" domain="0x0000" bus="0x01" slot="0x00" function="0x{i}" {multi} />
      </interface>"""


  with open(os.path.join(xmlout, f"{hostname}.xml"), "wt") as of:
    of.write(
      libvirt.format(
        hostname=hostname,
        ram=ram,
        interfacecfg=interfacecfg,
        cloudinitiso=cloudinitiso,
        hostdrive=hostdrive
      )
    )


def main():

  if len(sys.argv) < 3:
    print("Usage: generate.py [stage] [hostsfile] [image source] [output]")
    return

  stage = sys.argv[1]
  hostsfile = sys.argv[2]
  img_src = sys.argv[3]
  outdir = sys.argv[4]

  if stage not in ['all', 'img', 'iso', 'xml']:
    print("Stage should be: all, img, iso or xml")
    return

  with open(hostsfile, 'rt') as f:

    for line in f:
      line = line.replace('\n','').strip()
      if not line or line.startswith('#'):
        # skip empty and comment lines
        continue

      hostname, ram, net, ipaddr = line.split(' ')
      ram=int(ram)
      ipaddr = None if ipaddr == '-' else ipaddr

      print(f"Working on {hostname}...")
      if stage in ['iso', 'all']:
        run_iso(outdir, hostname, ipaddr)

      if stage in ['img', 'all']:
        run_img(outdir,hostname, img_src)

      if stage in ['xml', 'all']:
        run_xml(outdir,hostname, ram, net, ipaddr)


if __name__ == "__main__":
  main()