#!/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)

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):
					# 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

					try:
						newConf.set(section, key, value)
					# Add a new section where needed
					except configparser.NoSectionError as e:
						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"]))

dlib_archive = '/tmp/dlib_latest.tar.gz'

log('Downloading dlib')

loader = which('curl')
LOADER_CMD = None
if loader:
	LOADER_CMD = [loader, '--silent', '--retry', '5', '--location', '--output']
else:
	loader = which('wget')
	LOADER_CMD = [loader, '--quiet', '--tries', '5', '--output-document']

cmd = LOADER_CMD + [dlib_archive, 'https://api.github.com/repos/davisking/dlib/tarball/latest']

handleStatus(sc(cmd))

DLIB_DIR = None

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))'
)
with tarfile.open(dlib_archive) as tf:
	for item in tf:
		# tarball contains directory davisking-dlib-<commit id>, so peek into archive for the name
		if not DLIB_DIR:
			DLIB_DIR = item.name

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

os.unlink(dlib_archive)

log("Building dlib")

# Start the build
cmd = ["python3", "setup.py", "install"]

flags = ''
with open('/proc/cpuinfo') as info:
	for line in info:
		if 'flags' in line:
			flags = line
			break

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"]

sp = subprocess.run(cmd, cwd=DLIB_DIR, stderr=subprocess.STDOUT)
handleStatus(sp.returncode)

# simple check for CUDA
cuda_used = 'DLIB WILL USE CUDA' in sp.stdout

log("Cleaning up dlib")

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

log("Temporary dlib files removed")


log("Configuring howdy")

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

# Will be filled with the actual output lines
outlines = []
# Will be fillled with lines that contain coloring
printlines = []
# Track if the new lines have been insterted yet
inserted = False

# Open the PAM config file
with open("/etc/pam.d/common-auth") as fp:
	for line in fp.readline():
		# Add the line to the output directly, we're not deleting anything
		outlines.append(line)

		# Print the comments in gray and don't insert into comments
		if line[:1] == "#":
			printlines.append("\033[37m" + line + "\033[0m")
		else:
			printlines.append(line)

			# If it's not a comment and we haven't inserted yet
			if not inserted:
				# Set both the comment and the linking line
				line_comment = "# Howdy IR face recognition\n"
				line_link	= "auth	sufficient			pam_python.so /lib/security/howdy/pam.py\n\n"

				# Add them to the output without any markup
				outlines.append(line_comment)
				outlines.append(line_link)

				# Make the print orange to make it clear what's being added
				printlines.append("\033[33m" + line_comment + "\033[0m")
				printlines.append("\033[33m" + line_link + "\033[0m")

				# Mark as inserted
				inserted = True

# Print a file Header
print("\033[33m" + ">>> START OF /etc/pam.d/common-auth" + "\033[0m")

# Loop though all printing lines and use the enters from the file
for line in printlines:
	print(line, end="")

# Print a footer
print("\033[33m" + ">>> END OF /etc/pam.d/common-auth" + "\033[0m" + "\n")

# Do not prompt for a yes if we're in no promt mode
if "HOWDY_NO_PROMPT" not in os.environ:
	# Ask the user if this change is okay
	print("Lines will be insterted in /etc/pam.d/common-auth as shown above")
	ans = input("Apply this change? [y/N]: ").strip().lower()

	# Abort the whole thing if it's not
	if ans not in ("y", "yes"):
		print("Interpreting as a \"NO\", aborting")
		sys.exit(1)

print("Adding lines to PAM\n")

# Write to PAM
with open("/etc/pam.d/common-auth", "w") as common_auth:
	common_auth.write("".join(outlines))

# Sign off
print("Installation complete.")