#!/usr/bin/python3
# Installation script to install howdy
# Executed after primary apt install


def col(id):
	"""Add color escape sequences"""
	if id == 1: return "\033[32m"
	if id == 2: return "\033[33m"
	if id == 3: return "\033[31m"
	return "\033[0m"


# Import required modules
import fileinput
import subprocess
import sys
import os
import re
import tarfile
from shutil import rmtree, which

# Don't run unless we need to configure the install
# Will also happen on upgrade but we will catch that later on
if "configure" not in sys.argv:
	sys.exit(0)


def log(text):
	"""Print a nicely formatted line to stdout"""
	print("\n>>> " + col(1) + text + col(0) + "\n")


def handleStatus(status):
	"""Abort if a command fails"""
	if (status != 0):
		print(col(3) + "Error while running last command" + col(0))
		sys.exit(1)


# Create shorthand for subprocess creation
sc = subprocess.call

# We're not in fresh configuration mode so don't continue the setup
if not os.path.exists("/tmp/howdy_picked_device"):
	# Check if we have an older config we can restore
	if len(sys.argv) > 2:
		if os.path.exists("/tmp/howdy_config_backup_v" + sys.argv[2] + ".ini"):
			# Get the config parser
			import configparser

			# Load th old and new config files
			oldConf = configparser.ConfigParser()
			oldConf.read("/tmp/howdy_config_backup_v" + sys.argv[2] + ".ini")
			newConf = configparser.ConfigParser()
			newConf.read("/lib/security/howdy/config.ini")

			# Go through every setting in the old config and apply it to the new file
			for section in oldConf.sections():
				for (key, value) in oldConf.items(section):

					# MIGRATION 2.3.1 -> 2.4.0
					# If config is still using the old device_id parameter, convert it to a path
					if key == "device_id":
						key = "device_path"
						value = "/dev/video" + value

					# MIGRATION 2.4.0 -> 2.5.0
					# Finally correct typo in "timout" config value
					if key == "timout":
						key = "timeout"

					try:
						newConf.set(section, key, value)
					# Add a new section where needed
					except configparser.NoSectionError:
						newConf.add_section(section)
						newConf.set(section, key, value)

			# Write it all to file
			with open("/lib/security/howdy/config.ini", "w") as configfile:
				newConf.write(configfile)

	sys.exit(0)

# Open the temporary file containing the device ID
in_file = open("/tmp/howdy_picked_device", "r")
# Load it in, it should be a string
picked = in_file.read()
in_file.close()

# Remove the temporary file
os.unlink("/tmp/howdy_picked_device")

log("Upgrading pip to the latest version")

# Update pip
handleStatus(sc(["pip3", "install", "--upgrade", "pip"]))

log("Downloading and unpacking data files")

# Run the bash script to download and unpack the .dat files needed
handleStatus(subprocess.call(["./install.sh"], shell=True, cwd="/lib/security/howdy/dlib-data"))

log("Downloading dlib")


if "HOWDY_SKIP_DLIB" in os.environ:
	print(col(2) + "SKIPPING DLIB INSTALLATION BECAUSE OF HOWDY_SKIP_DLIB FLAG" + col(0))
else:
	dlib_archive = "/tmp/v19.16.tar.gz"
	loader = which("wget")
	LOADER_CMD = None

	# If wget is installed, use that as the downloader
	if loader:
		LOADER_CMD = [loader, "--tries", "5", "--output-document"]
	# Otherwise, fall back on curl
	else:
		loader = which("curl")
		LOADER_CMD = [loader, "--retry", "5", "--location", "--output"]

	# Assemble and execute the download command
	cmd = LOADER_CMD + [dlib_archive, "https://github.com/davisking/dlib/archive/v19.16.tar.gz"]
	handleStatus(sc(cmd))

	# The folder containing the dlib source
	DLIB_DIR = None
	# A regex of all files to ignore while unpacking the archive
	excludes = re.compile(
		"davisking-dlib-\w+/(dlib/(http_client|java|matlab|test/)|"
		"(docs|examples|python_examples)|"
		"tools/(archive|convert_dlib_nets_to_caffe|htmlify|imglab|python/test|visual_studio_natvis))"
	)

	# Open the archive
	with tarfile.open(dlib_archive) as tf:
		for item in tf:
			# Set the destenation dir if unset
			if not DLIB_DIR:
				DLIB_DIR = "/tmp/" + item.name

			# extract only files sufficient for building
			if not excludes.match(item.name):
				tf.extract(item, "/tmp")

	# Delete the downloaded archive
	os.unlink(dlib_archive)

log("Building dlib")

cmd = ["sudo", "python3", "setup.py", "install"]
cuda_used = False
flags = ""

if "HOWDY_SKIP_DLIB" not in os.environ:
	# Get the CPU details
	with open("/proc/cpuinfo") as info:
		for line in info:
			if "flags" in line:
				flags = line
				break

	# Use the most efficient instruction set the CPU supports
	if "avx" in flags:
		cmd += ["--yes", "USE_AVX_INSTRUCTIONS"]
	elif "sse4" in flags:
		cmd += ["--yes", "USE_SSE4_INSTRUCTIONS"]
	elif "sse3" in flags:
		cmd += ["--yes", "USE_SSE3_INSTRUCTIONS"]
	elif "sse2" in flags:
		cmd += ["--yes", "USE_SSE2_INSTRUCTIONS"]

	# Compile and link dlib
	try:
		sp = subprocess.Popen(cmd, cwd=DLIB_DIR, stdout=subprocess.PIPE)
	except subprocess.CalledProcessError:
		print("Error while building dlib")
		raise

	# Go through each line from stdout
	while sp.poll() is None:
		line = sp.stdout.readline().decode("utf-8")

		if "DLIB WILL USE CUDA" in line:
			cuda_used = True

		print(line, end="")

	log("Cleaning up dlib")

	# Remove the no longer needed git clone
	del sp
	rmtree(DLIB_DIR)

	print("Temporary dlib files removed")

log("Installing OpenCV")

handleStatus(subprocess.call(["pip3", "install", "--no-cache-dir", "opencv-python"]))

log("Configuring howdy")

# Manually change the camera id to the one picked
for line in fileinput.input(["/lib/security/howdy/config.ini"], inplace=1):
	line = line.replace("device_path = none", "device_path = " + picked)
	line = line.replace("use_cnn = false", "use_cnn = " + str(cuda_used).lower())

	print(line, end="")

print("Camera ID saved")

# Secure the howdy folder
handleStatus(sc(["chmod 744 -R /lib/security/howdy/"], shell=True))

# Allow anyone to execute the python CLI
os.chmod("/lib/security/howdy", 0o755)
os.chmod("/lib/security/howdy/cli.py", 0o755)
handleStatus(sc(["chmod 755 -R /lib/security/howdy/cli"], shell=True))
print("Permissions set")

# Make the CLI executable as howdy
os.symlink("/lib/security/howdy/cli.py", "/usr/local/bin/howdy")
os.chmod("/usr/local/bin/howdy", 0o755)
print("Howdy command installed")

log("Adding howdy as PAM module")

# Activate the pam-config file
handleStatus(subprocess.call(["pam-auth-update --package"], shell=True))

# Sign off
print("Installation complete.")
