#!/usr/bin/env python3

try:
	import argon2
except:
	exit()

import sys
import json
import os
import time
import stat

max_cache_age = 18000
cache_file = "/tmp/jupyter_password"
homedir = "/home/"
jupyter_config_json = [".jupyter/jupyter_notebook_config.json", ".jupyter/jupyter_config.json"]
jupyter_config_py = [".jupyter/jupyter_notebook_config.py", ".jupyter/jupyter_config.py"]

def on_storage():
	role_file = "/etc/facter/facts.d/puppet.yaml"
	if not os.path.isfile(role_file):
		return False
	try:
		f = open(role_file, "r")
		lines = f.readlines()
		for l in lines:
			l = l.strip()
			if "role: storage" in l:
				return True
	except:
		return False
	return False

def test_jupyter_password_is_OK_json(file):
	try:
		a = argon2.PasswordHasher()
		with open(file, "r") as f:
			data = json.load(f)
			return a.verify(data["NotebookApp"]["password"].split(":")[1],"")
	except:
		return False
	return False

def test_jupyter_password_is_OK_py(file):
	try:
		a = argon2.PasswordHasher()
		f = open(file, 'r')
		lines = f.readlines()
		f.close()

		for line in lines:
			line = line.strip()
			if line.startswith("#") or len(line) == 0:
				continue
			l = line.split("=")
			if l[0].strip() == "c.NotebookApp.password":
				pwd = l[1].strip()
				if len(pwd) > 1 and pwd.startswith('"') and pwd.endswith('"') or pwd.startswith("'") and pwd.endswith("'"):
					pwd = pwd[1:-1]
				return a.verify(pwd,"")
	except:
		return False
	return False

def get_cache(file):
	if not os.path.isfile(file):
		return ""

	age = time.time() - os.stat(file)[stat.ST_MTIME]
	if age > max_cache_age:
		return ""

	try:
		f = open(file, "r")
		lines = f.readlines()
		f.close()
		result = ""
		for l in lines:
			result += l
		return result.strip()
	except:
		return ""
		
if not on_storage():
	exit()

print("<<<jupyter_password>>>")

result = get_cache(cache_file)
if len(result) > 0:
	print(result)
	exit()

try:
	cache = open(cache_file, "w")
	counter = 0
	for item in os.listdir(homedir):
		for jc in jupyter_config_json:
			userjc = os.path.join(homedir, item, jc)
			if os.path.isfile(userjc):
				rc = test_jupyter_password_is_OK_json(userjc)
				if rc:
					counter += 1
					cache.write(f"{item}\n")

		for jc in jupyter_config_py:
			userjc = os.path.join(homedir, item, jc)
			if os.path.isfile(userjc):
				rc = test_jupyter_password_is_OK_py(userjc)
				if rc:
					counter += 1
					cache.write(f"{item}\n")
	if counter == 0:
		cache.write("OK\n")
			
	cache.close()
except Exception as err:
	print(f"UNKNOWN - {err}")
	exit()
	
result = get_cache(cache_file)
if len(result) > 0:
	print(result)
else:
	print("UNKNOWN - empty output")
