import sys
import tempfile
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"
"""

netplan = """version: 2
ethernets:
  enp1s0:
     dhcp4: false
     dhcp6: false
     addresses: [ {addr}/24 ]
     gateway4: {gateway}
     nameservers:
       addresses: [ 8.8.8.8, 8.8.4.4 ]
       search: [ local ]
"""

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.0">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="block" device="disk">
      <driver name="qemu" type="raw" cache="none" io="native"/>
      <backingStore/>
      <source dev="{hostdrive}"/>
      <alias name="virtio-disk0"/>
      <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(hostname:str, ipaddr:str):
  isoout = "/home/marcsello/ultrabirbcluster/vm_generator_luna/iso"
  os.makedirs(isoout, exist_ok=True)
  
  cloudinit_fhandle, cloudinit_fname = tempfile.mkstemp(suffix='.yaml')
  with open(cloudinit_fhandle, "wt") as f:
	  f.write(
      cloudinit.format(hostname=hostname)
    )

  netplan_fname = None
  netplan_generated = False
  if ipaddr:
    if not ipaddr.startswith('@'):
      gateway = '192.168.' + ipaddr.split('.')[2] + '.1'  # Ezen Kristóf jót nevetett
      netplan_fhandle, netplan_fname = tempfile.mkstemp(suffix='.yaml')
      with open(netplan_fhandle, "wt") as f:
        f.write(
          netplan.format(
            addr=ipaddr,
            gateway=gateway,
          )
        )
      netplan_generated = True
    else:
      fname = ipaddr[1:]
      netplan_fname = fname


    if netplan_fname:
      os.system(f'cloud-localds -v {isoout}/cloudinit-{hostname}.iso --network-config={netplan_fname} {cloudinit_fname}')
    else:
      os.system(f'cloud-localds -v {isoout}/cloudinit-{hostname}.iso {cloudinit_fname}')
  
  os.unlink(cloudinit_fname)
  if netplan_generated:
    os.unlink(netplan_fname)
  

def run_img(hostname:str, img_src:str):
  lv_name = hostname.replace('-', '_')
  hostdrive = f"/dev/vmdatastore/{lv_name}"
  os.system(f'dd if={img_src} of={hostdrive} bs=4M')


def run_xml(outdir:str, hostname:str, ram:int, net:str):
  xmlout = os.path.join(outdir, 'xml')
  isoout = "/var/lib/libvirt/images/cloudiso"
  cloudinitiso = os.path.join(isoout, f"cloudinit-{hostname}.iso")
  lv_name = hostname.replace('-', '_')
  hostdrive = f"/dev/vmdatastore/{lv_name}"
  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 run_ssh(outdir:str, hostname:str, ipaddr:str):
  if not ipaddr:
    return

  sshout_file = os.path.join(outdir, 'ssh_config')
  with open(sshout_file, 'at') as f:
    f.write("\n".join([
      f"Host {hostname}",
      f"\tHostName {ipaddr}",
      f"\tProxyJump luna",
      "",""
    ]))
    

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', 'ssh']:
    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(hostname, ipaddr)

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

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

      if stage in ['ssh', 'all']:
        run_ssh(outdir, hostname, ipaddr)

if __name__ == "__main__":
  main()