'''
########################
#  SETUP 
########################
 > pip3 install trilium-py

 - get trilium token in trilium options->ETAPI and "create new token"
 - get trilium_htb_folder click on note to become folder -> "note info". Note ID is there.
 - set the folder's owned attributes to:
 #label:user=promoted,single,text #label:root=promoted,single,text #label:respect=promoted,single,text #user=0 #root=0 #respect=0 

 - create a note named "HTBMachineTemplate" and set it's "owned attributes" to:
 #template #label:User=promoted,single,text #label:Root=promoted,single,text #label:Tags=promoted,single,text 

 - get this page's ID and put in "trilium_template_id"
 - to get "in progress" todo colour - add to trilium demo/scripting/taskmanager/implementation/CSS:

 span.fancytree-node.inprogress .fancytree-title {
    color: orange !important;
 }
'''

########################
#  IGNORE THESE
########################

import requests
from bs4 import BeautifulSoup
from datetime import datetime, timezone
from htb_client import HTBClient
from trilium_py.client import ETAPI

########################
#  EDIT BELOW HERE
########################

htb_code = ''
trilium_server_url = 'https://notes.place.com'
trilium_token = ''
trilium_htb_folder = ''
trilium_template_id = ''

########################
#  LEAVE BELOW HERE
########################

def generate_newpage(machine):
	no = "<span style=\"color:hsl(0,75%,60%);\">No</span>"
	yes = "<span style=\"color:hsl(120,75%,60%);\">Yes</span>"
	user_colour = no
	if machine['authUserInUserOwns']:
		user_colour = yes
	root_colour = no
	if machine['authUserInRootOwns']:
		root_colour = yes

	status = "Retired"
	if machine['active']:
		status = "Active"

	release_str = machine['release']
	release_date = datetime.strptime(release_str, "%Y-%m-%dT%H:%M:%S.%fZ")
	formatted_release_date = release_date.strftime("%d %B %Y")

	html = """
	<figure class="image image-style-align-left image_resized" style="width:12.28%;">
	  <img src="https://www.hackthebox.com{avatar}">
	</figure>
	<p>
	  <strong>OS:</strong> {os}<br>
	  <strong>Difficulty:</strong> {difficultyText} <br>
	  <strong>Rating:</strong> {rating} / 5<br>
	  <strong>Points:</strong> {points}<br>
	  <strong>User / Root: </strong> {user_colour} / {root_colour}<br>
	  <strong>Released:</strong> {release_date}<br>
	  <strong>State:</strong> {status}
	</p>
	<hr>
	<h2>Notes</h2>
	<p>&nbsp;</p>
	""".format(
		os=machine['os'],
		difficultyText=machine['difficultyText'],
		rating=machine['star'],
		points=machine['points'],
		release_date=formatted_release_date,
		user_colour=user_colour,
		root_colour=root_colour,
		status = status,
		avatar = machine['avatar'],
	)

	return html

def get_timestamp(machine):
    # Parse the release date string into a datetime object
    release_str = machine['release']
    dt = datetime.strptime(release_str, "%Y-%m-%dT%H:%M:%S.%fZ")
    
    # Set the timezone to UTC
    dt = dt.replace(tzinfo=timezone.utc)
    
    # Get the timestamp
    timestamp = dt.timestamp()
    return timestamp

print("[+] connecting to HTB")
client = HTBClient(password=htb_code)
print("[+] connecting to trilium")
ea = ETAPI(trilium_server_url, trilium_token)
print("[i] version: ", ea.app_info()['appVersion'])

print("[i] HTB User:", client.user['id'], "-", client.user['name'])
print("[i] user owns:", client.user['user_owns'], "| Root owns:", client.user['root_owns'], "| Respect:", client.user['respects'])

master_folder = ea.get_note(trilium_htb_folder)
for attribute in master_folder['attributes']:
	if attribute['name'] == "user":
		if attribute['value'] != str(client.user['user_owns']):
			print("[+] updating user owns (folder attribute)")
			ea.patch_attribute(attributeId=attribute['attributeId'], value=str(client.user['user_owns']))
	if attribute['name'] == "root":
		if attribute['value'] != str(client.user['root_owns']):
			print("[+] updating root owns (folder attribute)")
			ea.patch_attribute(attributeId=attribute['attributeId'], value=str(client.user['root_owns']))
	if attribute['name'] == "respect":
		if attribute['value'] != str(client.user['respects']):
			print("[+] updating respect (folder attribute)")
			ea.patch_attribute(attributeId=attribute['attributeId'], value=str(client.user['respects']))

print("[+] gathering machines info")
machines = client.get_all_machines()
machines.sort(key=get_timestamp)

print(f"[i] Retrieved {len(machines)} machines")
#for machine in machines:
#    print(f" - ID: {machine['id']}, Name: {machine['name']}, OS: {machine['os']}, Difficulty: {machine['difficultyText']}")

machine_count = 0
completed_count = 0
for machine in machines:
	machine_count += 1
	print('processing: ',machine_count, "/", len(machines), "("+machine['name']+")               " , end='\r')

	res = ea.search_note(
		search="\""+machine['name']+"\"",
		ancestorNoteId=trilium_htb_folder,
		ancestorDepth='eq1',
		limit=1,
		fastSearch=True,
	)

	if res['results'] and machine['name'] == res['results'][0]['title']:
		# page exists - lets check if the details have changed

		current_html = ea.get_note_content(noteId=res['results'][0]['noteId'])
		current_soup = BeautifulSoup(current_html, 'html.parser')
		current_paragraph = current_soup.find_all('p')[0].text

		new_html = generate_newpage(machine)
		new_soup = BeautifulSoup(new_html, 'html.parser')
		new_paragraph = new_soup.find_all('p')[0].text

		# current page contains first paragraph of "blank" (useful for when it doesnt create or find the note properly.. shouldnt get here)
		if current_paragraph == "blank":
			ea.update_note_content(noteId=res['results'][0]['noteId'], content=new_html)
			ea.create_attribute(attributeId=None, isInheritable=False, noteId=res['results'][0]['noteId'], type="relation", name="template", value=trilium_template_id)
			if machine['authUserInUserOwns'] and machine['authUserInRootOwns']:
				ea.create_attribute(attributeId=None, isInheritable=False, noteId=res['results'][0]['noteId'], type="label", name="cssClass", value="done")
			else:
				if machine['authUserInUserOwns'] or machine['authUserInRootOwns']:
					ea.create_attribute(attributeId=None, isInheritable=False, noteId=res['results'][0]['noteId'], type="label", name="cssClass", value="inprogress")
				else:
					ea.create_attribute(attributeId=None, isInheritable=False, noteId=res['results'][0]['noteId'], type="label", name="cssClass", value="todo")
			# re-get the current content
			current_html = ea.get_note_content(noteId=res['results'][0]['noteId'])
			current_soup = BeautifulSoup(current_html, 'html.parser')

		if current_paragraph != new_paragraph:

			# details have updated!
			print("[+] updating page:",machine['name'], "-> "+res['results'][0]['title']+"                 ")
			replacement = current_soup.find('p')
			replacement.replace_with( new_soup.find_all('p')[0] )
			ea.update_note_content(noteId=res['results'][0]['noteId'], content=current_soup)

			# now to update the label
			for attribute in res['results'][0]['attributes']:
				if attribute['name'] == "cssClass":
					if machine['authUserInUserOwns'] and machine['authUserInRootOwns']:
						ea.patch_attribute(attributeId=attribute['attributeId'], value="done")
					else:
						if machine['authUserInUserOwns'] or machine['authUserInRootOwns']:
							ea.patch_attribute(attributeId=attribute['attributeId'], value="inprogress")
						else:
							ea.patch_attribute(attributeId=attribute['attributeId'], value="todo")
							
	else:
		# title does not exist - create the note
		html = generate_newpage(machine)
		new_note = ea.create_note(
			parentNoteId=trilium_htb_folder,
			type="text",
			title=machine['name'],
			content=html,
		)
		print("[+] created note:", machine['name'], "                        ")

		ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="relation", name="template", value=trilium_template_id)
		if machine['authUserInUserOwns'] and machine['authUserInRootOwns']:
			ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="cssClass", value="done")
		else:
			if machine['authUserInUserOwns'] or machine['authUserInRootOwns']:
				ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="cssClass", value="inprogress")
			else:
				ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="cssClass", value="todo")

	if machine['authUserInUserOwns'] and machine['authUserInRootOwns']:
			completed_count += 1

print("[+] updating folder name                                          ")
ea.patch_note(noteId=trilium_htb_folder,title="Machines - "+str(completed_count)+" / "+str(len(machines)))

print("[=] processed", machine_count)