#!/usr/bin/env python3

import argparse
import glob
import os
import stat
import urllib.request
import zipfile
import tempfile
import re
# import errno
from dataclasses import dataclass
from shutil import copy2 as cp, which
from subprocess import call, Popen
from typing import Optional


class FatalError(Exception):
	def __init__(self, message, return_code=None):
		super().__init__(message)
		self.code = return_code


def get_img_file(printable_name: str, destination_file: str, url: str, filename_in_zip: str) -> None:
	"""
	Download a zip file from the internet, extract a single file and place it somewhere.
	If filename_in_zip is the empty string, the downloaded file will be renamed/moved instead.

	:param printable_name: description of what this thing is, used to generate temporary file names and to show in messages
	:param destination_file: name of the extracted file, will be placed in tftp/boot/img/included/
	:param url: where to download from
	:param filename_in_zip: This will be extracted and moved to destination_file.
	:return:
	"""
	if os.path.isfile("tftp/boot/img/included/{0}".format(destination_file)):
		print("{0} found".format(printable_name))
	else:
		if filename_in_zip == "":
			# Not a zip file
			print("Getting {0}...".format(printable_name))
			urllib.request.urlretrieve(url, filename="tftp/boot/img/included/{0}".format(destination_file))
		else:
			# Zip file: download and extract
			temp_name = "pxenow-downloaded-{0}.zip".format(printable_name)

			if not os.path.isfile(temp_name):
				print("Getting {0}...".format(printable_name))
				urllib.request.urlretrieve(url, filename=temp_name)

			with zipfile.ZipFile(temp_name, "r") as zip_ref:
				print("Extracting {0}...".format(filename_in_zip))
				zipinfo = zip_ref.getinfo(filename_in_zip)
				zipinfo.filename = destination_file  # avoids extracting random folders and renaming
				zip_ref.extract(filename_in_zip, path="tftp/boot/img/included/")
			os.remove(temp_name)

	return


def guess_interface() -> str:
	"""
	Guess which network interface to use (= pick the only one other than lo and exit if more than one exists)

	:return: interface, e.g. enp3s0
	"""
	import netifaces

	interfaces = netifaces.interfaces()

	if 'lo' in interfaces:
		interfaces.remove('lo')

	num = len(interfaces)
	if num == 0:
		print("No interfaces found (other than loopback, possibly)")
		exit(1)
	elif num == 1:
		return interfaces[0]
	else:
		print("More than one interface: specify which one to use in " + str(interfaces))
		exit(1)


def relink(source: str, destination: str) -> None:
	"""
	Create a symlink. If it already exists, make sure it is correct.

	:param source: real file (absolute path required!)
	:param destination: the link itself
	:return: None
	"""

	if os.path.islink(destination):
		current_source = os.readlink(destination)
		if current_source == source:
			return
		else:
			os.remove(destination)
	elif os.path.exists(destination):
		raise FatalError("Cannot create link from {}: file already exists and is not a link".format(destination))

	print("Creating symlink {frm} pointing to {to}...".format(frm=destination, to=source))
	os.symlink(source, destination)


def configure_nfs(exports_file_content: str, use_sudo: bool) -> None:
	print("Stopping NFS server...")
	sudocall(["systemctl", "stop", "nfs-server"])

	different = False

	try:
		with open(ETCEXPORTS, "r") as exports_file:
			old_exports = exports_file.read()
			if old_exports != exports_file_content:
				different = True
	except IOError as why:
		raise FatalError("Cannot read {}: {}".format(ETCEXPORTS, str(why)))

	if different:
		print("Saving new {0}...".format(ETCEXPORTS))

		try:
			# Try to write new file
			with open(ETCEXPORTS, "w") as exports_file:
				exports_file.write(exports_file_content)
		except PermissionError as why:
			# Cannot write /etc/exports
			if use_sudo:
				# Create a temporary file and overwrite /etc/exports (this allows calling sudo)
				write_content_sudocall(exports_file_content, ETCEXPORTS)
			else:
				raise FatalError("Cannot write {}: {}".format(ETCEXPORTS, str(why)))

		print("Reloading exports...")
		sudocall(["exportfs", "-a"])
	else:
		print("No need to update {0}".format(ETCEXPORTS))

	print("Starting NFS server...")
	sudocall(["systemctl", "start", "nfs-server"])
	print("")


def sudocall_wait(what: list) -> int:
	if args.sudo:
		callme = ["sudo"]
		callme.extend(what)
	else:
		callme = what

	with Popen(callme) as process:
		try:
			return process.wait()
		except KeyboardInterrupt:
			raise
		except Exception:
			process.kill()
			process.wait()
			raise


def sudocall(what: list) -> int:
	"""
	Call some command, prepending sudo if args.sudo is True

	:param what: command and arguments
	:return: whatever sudo/the command returns
	"""
	if args.sudo:
		callme = ["sudo"]
		callme.extend(what)
	else:
		callme = what

	return call(callme)


def write_content_sudocall(content: str, destination: str) -> None:
	"""
	Write content to a file with a command called via sudocall.
	Useful to write files where current user has no write permission.

	:param content: File content
	:param destination: Where to place it
	:return: Nothing
	"""
	with tempfile.NamedTemporaryFile() as tmp_file, open(tmp_file.name, 'w') as tmp:
		try:
			tmp.write(content)
			tmp.flush()
		except IOError:
			print("Cannot write to temporary file {0}".format(tmp_file.name))
			return
		sudocall_result = sudocall(["cp", tmp_file.name, destination])
		if sudocall_result != 0:
			raise FatalError("Cannot call cp {} {}".format(tmp_file.name, destination), sudocall_result)


@dataclass
class EntryParameters:
	name: str
	server: str
	vmlinuz: str
	initrd: str
	root_directory: str
	keymap: Optional[tuple]


def get_syslinux_config_for(name: str, root_directory: str, server: str, keymap: Optional[tuple]) -> str:
	"""
	Generate a Syslinux/PXELINUX configuration entry for an ISO/directory

	:param name: Entry name (human-readable)
	:param server: Server IP
	:param root_directory: Directory where ISO was mounted, or extracted
	:param keymap: Keymaps, for entries that support them
	:return:
	"""

	maybe_vmlinuz = (
		f"{root_directory}/casper/vmlinuz*",
		# f"{root_directory}/casper/vmlinuz.efi",
		f"{root_directory}/arch/boot/x86_64/vmlinuz",
		f"{root_directory}/live/vmlinuz-*",
	)
	maybe_initrd = (
		f"{root_directory}/casper/initrd.lz",
		# f"{root_directory}/casper/initrd",
		f"{root_directory}/arch/boot/x86_64/archiso.img",
		f"{root_directory}/live/initrd.img*"
	)

	vmlinuz_src = search_for("vmlinuz", maybe_vmlinuz, name)
	initrd_src = search_for("initrd", maybe_initrd, name)

	vmlinuz = "mounted/{0}-vmlinuz".format(name)
	initrd = "mounted/{0}-initrd".format(name)

	relink(vmlinuz_src, "tftp/boot/" + vmlinuz)
	relink(initrd_src, "tftp/boot/" + initrd)

	ep = EntryParameters(
		name=name,
		server=server,
		vmlinuz=vmlinuz,
		initrd=initrd,
		root_directory=root_directory,
		keymap=keymap
	)

	if os.path.isdir(f"{root_directory}/arch/boot/x86_64/"):
		print(f"{name} seems to be Arch Linux x86_64")
		return syslinux_config_arch(ep)
	elif os.path.isdir(f"{root_directory}/live/") and (os.path.isdir(f"{root_directory}/d-i/") or os.path.exists(f"{root_directory}/DEBIAN_CUSTOM")):
		print(f"{name} seems to be Debian live")
		return syslinux_config_debian(ep)
	else:
		print(f"{name} seems to be Ubuntu (derivative) or unrecognized")
		return syslinux_config_ubuntu(ep)


def name_compact(name: str):
	return re.sub(r'\W+', '', name)


def syslinux_config_debian(ep: EntryParameters):
	return str(f"""
	LABEL {name_compact(ep.name)}
	MENU LABEL {ep.name}
	TEXT HELP
		Boot Debian with kernel+initramfs from TFTP,
		and root fs from NFS server {ep.server}
	ENDTEXT
	LINUX {ep.vmlinuz}
	APPEND root=/dev/nfs netboot=nfs nfsroot={ep.server}:{ep.root_directory}/ initrd=::/boot/{ep.initrd} boot=live
	""")


# localized debian live:
# LABEL Italian (it)
#   SAY "Booting Italian (it)..."
#   linux /live/vmlinuz-4.9.0-7-amd64
#   APPEND initrd=/live/initrd.img-4.9.0-7-amd64 boot=live components locales=it_IT.UTF-8
#
# Somewhat recent pxe boot parameters:
# https://lists.debian.org/debian-live/2016/01/msg00008.html


def syslinux_config_ubuntu(ep: EntryParameters):
	label = name_compact(ep.name)
	config = ""

	if ep.keymap is None:
		keymaps = [""]
	else:
		keymaps = ep.keymap

	for k in keymaps:
		if k == "":
			keyboard_description = ""
			keyboard_param = ""
		else:
			keyboard_description = ", {} layout".format(k)
			keyboard_param = " keyboard-configuration/layoutcode?=" + k
		config += f"""
	LABEL {label}{k}
	MENU LABEL {ep.name} (NFS{keyboard_description})
	TEXT HELP
		Boot with kernel+initrd from TFTP,
		and squashfs from NFS server {ep.server}
	ENDTEXT
	KERNEL {ep.vmlinuz}
	APPEND root=/dev/nfs boot=casper netboot=nfs nfsroot={ep.server}:{ep.root_directory} initrd={ep.initrd}{keyboard_param} ---
		"""

	return config


# TODO: doesn't work
def syslinux_config_arch(ep: EntryParameters):
	return str(f"""
	LABEL {name_compact(ep.name)}
	MENU LABEL {ep.name}
	TEXT HELP
		Boot Arch Linux with kernel+initramfs from TFTP,
		and root fs from NFS server {ep.server}
	ENDTEXT
	LINUX {ep.vmlinuz}
	INITRD {ep.initrd}
	APPEND ip=:: archisobasedir=arch archiso_nfs_srv={ep.server}:{ep.root_directory}
	SYSAPPEND 2
	""")


def search_for(what: str, where: tuple, iso_name: str) -> str:
	"""
	Find first file that exists from a list of paths

	:param what: File name/description (generic, human-readable)
	:param where: Possible paths (to be globbed)
	:param iso_name: ISO name or other printable name
	:return:
	"""
	for maybe in where:
		matches = glob.glob(maybe)
		if len(matches) > 0:
			maybe = matches[0]
			if len(matches) > 1:
				print(f"Warning: ambiguous globbing for {what}, choosing {maybe}")
			return maybe
	raise FatalError(f"{what} not found in {iso_name}")


def get_network_parameters(iface: Optional[str], ip: Optional[str], nmask: Optional[str]) -> {Optional[str], str, str}:
	"""
	Obtain at least an IP and a netmask from supplied parameters.

	:param iface: Interface
	:param ip: IP
	:param nmask: Netmask
	:return: Interface, IP, netmask
	"""
	if ip is None:
		try:
			import netifaces
		except ModuleNotFoundError:
			raise FatalError("Module netifaces not found, install it or specify IP and netmask (-s and -i parameters)")

		# No IP given and no interface, so we have to guess interface and obtain IP
		if iface is None:
			iface = guess_interface()

		addresses = netifaces.ifaddresses(iface)[netifaces.AF_INET]

		if len(addresses) == 0:
			raise FatalError("No addresses found for interface " + str(args.interface))

		ip = addresses[0]['addr']
		nmask = addresses[0]['netmask']
	else:
		if nmask is None:
			raise FatalError("Provide a netmask if manually setting the IP, e.g. -n 255.255.255.0")
		ip = args.server

	return iface, ip, nmask


def copy_files_paths(destination: str, filename: str, source_paths: tuple):
	if len(source_paths) < 1:
		raise Exception(f"Trying to copy {filename} to {destination} but no source paths provided")

	found = False
	for path in source_paths:
		src = path + filename
		if os.path.isfile(src):
			found = True
			break
	if found:
		print(f"Copying {filename} (found at {src})")
		# noinspection PyUnboundLocalVariable
		cp(src, destination)
	else:
		raise FatalError(f"Cannot find {filename}, tried: {str(source_paths)}")


def copy_pxelinux_files() -> None:
	# First path is Arch Linux, second is (X)ubuntu
	pxelinux_dirs = ("/usr/lib/syslinux/bios/", "/usr/lib/PXELINUX/")
	syslinux_modules_dirs = ("/usr/lib/syslinux/bios/", "/usr/lib/syslinux/modules/bios/")
	syslinux_other_dirs = ("/usr/lib/syslinux/bios/", "/usr/lib/syslinux/")
	pci_ids_path = ("/usr/share/hwdata/", "/usr/share/misc/")

	files = {
		"pxelinux.0": pxelinux_dirs,
		"lpxelinux.0": pxelinux_dirs,
		"ldlinux.c32": syslinux_modules_dirs,
		"vesamenu.c32": syslinux_modules_dirs,
		"menu.c32": syslinux_modules_dirs,
		"libutil.c32": syslinux_modules_dirs,
		"hdt.c32": syslinux_modules_dirs,  # for HDT
		"libmenu.c32": syslinux_modules_dirs,  # for HDT
		"libcom32.c32": syslinux_modules_dirs,  # for HDT
		"libgpl.c32": syslinux_modules_dirs,  # for HDT
		"reboot.c32": syslinux_modules_dirs,
		"poweroff.c32": syslinux_modules_dirs,
		"memdisk": syslinux_other_dirs,
		"pci.ids": pci_ids_path
	}

	for filename, paths in files.items():
		try:
			copy_files_paths("tftp/boot/" + filename, filename, paths)
		except (FileNotFoundError, OSError) as why:
			print()
			raise FatalError(f"Cannot copy file: {str(why)}")


def generate_dnsmasq_configuration(interface: Optional[str], server: str) -> None:
	"""
	Generate, validate and save dnsmasq configuration

	:param interface: Network interface
	:param server: IP of the TFTP and NFS server
	:return: Nothing
	"""
	print("Generating dnsmasq configuration...")
	# useful stuff: https://wiki.archlinux.org/index.php/Dnsmasq#PXE_server
	dnsmasq_generated_config = ""

	# Could still be "none" when providing IP and netmask, but it's not mandatory
	if interface is not None:
		dnsmasq_generated_config += """interface = {iface}
	bind-interfaces
	""".format(iface=interface)

	dnsmasq_generated_config += """port=0
	dhcp-leasefile=dnsmasq.leases

	enable-tftp
	tftp-root={pwd}/tftp
	dhcp-option-force=66,{ip} # TFTP server

	log-dhcp # More info on what's going on
	dhcp-no-override # Don't put useless fields in DHCP offer
	dhcp-range={ip},proxy

	pxe-service=X86PC, "PXE boot NOW", boot/lpxelinux # It adds the .0 by itself
	dhcp-option=vendor:PXEClient,6,2b # "kill multicast"

	# These don't work in Proxy DHCP mode, apparently:
	# dhcp-boot=boot/lpxelinux.0
	# dhcp-option-force=209,"pxelinux.cfg/default" # PXELINUX config file (it's the default anyway)
	# dhcp-option-force=210,/boot/ # PathPrefix, doesn't seem to be needed (See RFC 5071)
	""".format(ip=server, pwd=os.getcwd())

	with open(DNSMASQCONF, "w") as config_file:
		config_file.write(dnsmasq_generated_config)
	os.chmod(DNSMASQCONF, os.stat(DNSMASQCONF).st_mode | stat.S_IROTH)

	if which("dnsmasq") is None:
		raise FatalError("dnsmasq binary not found in PATH")

	sudocall_result = sudocall(["dnsmasq", "-d", "--test", "-C", DNSMASQCONF])
	if sudocall_result != 0:
		raise FatalError("dnsmasq returned " + str(sudocall_result), sudocall_result)


def find_bootable_files(paths) -> {list, list, list}:
	"""
	Check that supplied files and directories exist and determine what they are

	:param paths: Path to ISO files, IMG files, directories and whatever else
	:return: ISO files, IMG files, and directories
	"""
	iso_files = []
	img_files = []
	iso_directories = []

	if paths is not None:
		for file in args.iso:
			if os.path.isdir(file):
				iso_directories.append(file)
			elif file.endswith(".iso"):
				iso_files.append(file)
			else:
				img_files.append(file)

	return iso_files, img_files, iso_directories


def mount_iso(path_to_iso: str, mounted_files: list) -> {str, str}:
	"""
	Mount an ISO file

	:param path_to_iso: ISO file
	:param mounted_files: Global list of mounted files
	:return: Filename and directory where it was mounted
	"""
	if not os.path.isfile(path_to_iso):
		raise FatalError("{0} doesn't exist".format(path_to_iso))

	file_name = os.path.basename(path_to_iso)
	mount_directory = "{cwd}/nfs/{name}.mount".format(cwd=os.getcwd(), name=file_name)

	if " " in mount_directory:
		raise FatalError(
			"NFS path {0} contains spaces so it won't work, move everything in a directory without spaces and retry"
			.format(mount_directory))

	os.makedirs(mount_directory, exist_ok=True)

	# Could be done "if result == 0", but if script crashes midway you'd need to unmount manually...
	mounted_files.append(mount_directory)

	# TODO: detect if already mounted?
	sudocall_result = sudocall(["mount", "-o", "loop,ro", "-t", "iso9660", path_to_iso, mount_directory])

	if sudocall_result != 0:
		raise FatalError("mount returned " + str(sudocall_result), sudocall_result)

	return file_name, mount_directory


def generate_iso_config(iso: list, keymaps: Optional[tuple], mounted_files: list, server: str):
	syslinux_generated_config = ""

	if len(iso) > 0:
		print("ISO files: {0}\n".format(str(iso)))
		os.makedirs("nfs", exist_ok=True)
		os.makedirs("tftp/boot/mounted", exist_ok=True)

		syslinux_generated_config = "\tMENU SEPARATOR\n\n"

	for file in iso:
		name, mount_directory = mount_iso(file, mounted_files)
		print("Generating syslinux config for {0}".format(name))
		syslinux_generated_config += get_syslinux_config_for(name, mount_directory, server, keymaps)

	with open(GENERATEDISOCONF, "w") as config_file:
		config_file.write(syslinux_generated_config)
	os.chmod(GENERATEDISOCONF, os.stat(GENERATEDISOCONF).st_mode | stat.S_IROTH)


def generate_isd_config(isd: list, keymaps: Optional[tuple], mounted_files: list, server: str):
	syslinux_generated_config = ""

	if len(isd) > 0:
		syslinux_generated_config = "\tMENU SEPARATOR\n\n"

	# TODO: implement

	with open(GENERATEDISDCONF, "w") as config_file:
		config_file.write(syslinux_generated_config)
	os.chmod(GENERATEDISDCONF, os.stat(GENERATEDISDCONF).st_mode | stat.S_IROTH)


def generate_img_config(img: list):
	syslinux_generated_config = ""

	if len(img) > 0:
		print("IMG files: {0}\n".format(str(img)))
		os.makedirs("tftp/boot/img", exist_ok=True)

		syslinux_generated_config = "\tMENU SEPARATOR\n\n"

	for file in img:
		file_name = os.path.basename(file)

		print("Generating syslinux config for {0}".format(img))
		syslinux_generated_config += """
	LABEL -
	MENU LABEL {name}
	TEXT HELP
	Launch {name} as image file
	ENDTEXT
	KERNEL memdisk
	APPEND initrd=img/{name}
		""".format(name=file_name)
		relink(file, "tftp/boot/img/" + file_name)

	with open(GENERATEDIMGCONF, "w") as config_file:
		config_file.write(syslinux_generated_config)
	os.chmod(GENERATEDIMGCONF, os.stat(GENERATEDIMGCONF).st_mode | stat.S_IROTH)


def main(mounted_files) -> int:
	# Quick setting a static IP:
	# IFACE=enp3s0
	# ip link set "${IFACE}" up
	# ip addr flush dev "${IFACE}"
	# ip addr add 10.80.7.1/24 dev "${IFACE}"
	# flush is necessary only if this is done over and over again (e.g. in a script)

	interface, server, netmask = get_network_parameters(args.interface, args.server, args.netmask)
	del args.interface
	del args.server

	if args.keymaps is None:
		keymaps = None
	else:
		keymaps = args.keymaps.split(",")
	del args.keymaps

	copy_pxelinux_files()

	os.makedirs("tftp/boot/iso", exist_ok=True)
	os.makedirs("tftp/boot/img/included", exist_ok=True)

	generate_dnsmasq_configuration(interface, server)

	get_img_file("Memtest86+", "memtest", "http://www.memtest.org/download/5.01/memtest86+-5.01.zip", "memtest86+-5.01.bin")
	get_img_file("PLoP", "plpbt.bin", "https://download.plop.at/files/bootmngr/plpbt-5.0.15.zip", "plpbt-5.0.15/plpbt.bin")
	get_img_file("netboot.xyz", "netbootxyz", "https://boot.netboot.xyz/ipxe/netboot.xyz.lkrn", "")

	# ISO files, IMG files, ISo Directory, just to keep the symmetry of three-letter codes
	iso, img, isd = find_bootable_files(args.iso)

	generate_iso_config(iso, keymaps, mounted_files, server)
	generate_isd_config(isd, keymaps, mounted_files, server)
	generate_img_config(img)

	exports = ""

	for full_export in mounted_files:
		exports += "{export} {ip}/{nm}(ro,all_squash,insecure,no_subtree_check)\n"\
			.format(export=full_export, ip=server, nm=netmask)

	if exports != "":
		if args.nfs:
			configure_nfs("# Created by pxenow\n" + exports, args.sudo)
		else:
			print("Almost done: add this to /etc/exports, run 'exportfs -a' and start your NFS server:")
			print("---------------------------------------------------------------------------------")
			print(exports)
			print("---------------------------------------------------------------------------------")

	try:
		returned = sudocall_wait(["dnsmasq", "-d", "-C", DNSMASQCONF])
		if returned != 0:
			raise FatalError("Failed launching dnsmasq", returned)
	except KeyboardInterrupt:
		print("Keyboard interrupt (ctrl+C) detected")

	if args.nfs:
		print("Stopping NFS server...")
		sudocall(["systemctl", "stop", "nfs-server"])

	return 0


if __name__ == "__main__":
	parser = argparse.ArgumentParser(description='Create a PXE server right here, right now.')
	parser.add_argument('iso', nargs='*', type=str, help="Path to ISO images")
	parser.add_argument('-i', '--interface', type=str, help="Interface to bind, e.g. enp3s0")
	parser.add_argument('-n', '--netmask', type=str, help="Netmask, used only if -s is also used")
	parser.add_argument('-s', '--server', type=str,
		help="IP address of current machine, used as TFTP, DHCP and NFS server")
	parser.add_argument('-k', '--keymaps', type=str, help="Comma-separated list of keymaps")
	# parser.add_argument('-m', '--memdisk', action='store_true', help="Generate memdisk entries for each ISO")
	parser.add_argument('-N', '--nfs', action='store_true', help="Blindly overwrite /etc/exports and manage NFS server")
	parser.add_argument('-S', '--sudo', action='store_true', help="Use sudo for commands that require root permissions")
	parser.set_defaults(memdisk=False)
	parser.set_defaults(nfs=False)
	parser.set_defaults(sudo=False)
	args = parser.parse_args()

	mounted = []
	DNSMASQCONF = "dnsmasq-pxe.conf"
	GENERATEDIMGCONF = "tftp/boot/pxelinux.cfg/img-generated.conf"
	GENERATEDISOCONF = "tftp/boot/pxelinux.cfg/iso-generated.conf"
	GENERATEDISDCONF = "tftp/boot/pxelinux.cfg/isd-generated.conf"
	ETCEXPORTS = "/etc/exports"

	mounted_list = []
	code = 1
	try:
		code = main(mounted_list)
	except FatalError as e:
		print(str(e))
		if e.code is None:
			code = 1
		else:
			code = e.code
	finally:
		if len(mounted_list) > 0:
			errors = False
			print("Unmounting all ISOs...")
			for directory in mounted_list:
				print("Unmounting {0}...".format(directory))
				result = sudocall(["umount", directory])
				if result == 0:
					try:
						os.rmdir(directory)
					except OSError:
						pass
				else:
					errors = True
					print("Failed to unmount {}: {}".format(directory, os.strerror(result)))
			if errors != 0:
				print("Warning: some umount(s) failed, try stopping/restarting the NFS server and re-running the script")
	exit(code)
