import os
import json
import requests
from collections import defaultdict
from bs4 import BeautifulSoup
from datetime import datetime, timezone
from htb_client import HTBClient
from trilium_py.client import ETAPI
 
# Get the absolute path of the script's directory
script_dir = os.path.dirname(os.path.abspath(__file__))
 
# Construct the full path to the config.json file
config_path = os.path.join(script_dir, 'config.json')
 
# Load configuration from the JSON file
with open(config_path, 'r') as f:
    config = json.load(f)
 
# Accessing config values
htb_code = config['htb_code']
trilium_server_url = config['trilium_server_url']
trilium_token = config['trilium_token']
trilium_challenges_folder = config['trilium_challenges_folder']
trilium_challenges_template_id = config['trilium_challenges_template_id']
 
def get_timestamp(machine):
    # Parse the release date string into a datetime object
    release_str = machine['release_date']
    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("[+] gathering challenges info")
categories = defaultdict(list)
challenges = client.get_all_challenges()
challenges.sort(key=get_timestamp)
total_completed = 0  # Variable to track total completed challenges
 
print(f"[i] Retrieved {len(challenges)} challenges")
 
# Group challenges by their categories
for challenge in challenges:
    categories[challenge['category_name']].append(challenge)
    if challenge['is_owned']:
        total_completed += 1  # Increment total completed if challenge is owned
 
# Print out the grouped challenges with the number of completed and total challenges in each category
for category, grouped_challenges in categories.items():
    total = len(grouped_challenges)
    completed = sum(1 for challenge in grouped_challenges if challenge['is_owned'])  # Count completed challenges
    
    res = ea.search_note(
        search=f"note.title %= '{category}*'",
        ancestorNoteId=trilium_challenges_folder,
        ancestorDepth='eq1',
        limit=1,
        fastSearch=True,
    )
    
    catId = ""
    if res['results'] and res['results'][0]['title'].split(' - ')[0].strip().lower() == category.lower():
        # page exists - lets check if the details have changed
        ea.patch_note(noteId=res['results'][0]['noteId'], title=category+" - "+str(completed)+" / "+str(total))
        catId = res['results'][0]['noteId']
        print(f"[i] updated category: {category} - ({completed}/{total})")
    else:
        new_note = ea.create_note(
            parentNoteId=trilium_challenges_folder,
            type="text",
            title=category+" - "+str(completed)+" / "+str(total),
            content=" ",
        )
        catId = new_note['note']['noteId']
        print(f"[+] created category: {catId} {category} - ({completed}/{total})")
    
    for challenge in grouped_challenges:
        #print(f" - ID: {challenge['id']}, Name: {challenge['name']}, Difficulty: {challenge['difficulty']}")
        res2 = ea.search_note(
            search=f"{challenge['name']}",
            ancestorNoteId=catId,
            ancestorDepth='eq1',
            limit=1,
            fastSearch=True,
        )
        # already exists update the values
        if res2['results'] and res2['results'][0]['title'].lower() == challenge['name'].lower():
            print(f"[i] found ID: {challenge['id']}, Name: {challenge['name']}, Difficulty: {challenge['difficulty']}")
            #print(f"Search response for challenge '{challenge['name']}': {res2}")
            for attribute in res2['results'][0]['attributes']:
                if attribute['name'] == "Difficulty":
                    ea.patch_attribute(attributeId=attribute['attributeId'], value=challenge['difficulty'])
                if attribute['name'] == "Released":
                    release_str = challenge['release_date']
                    release_date = datetime.strptime(release_str, "%Y-%m-%dT%H:%M:%S.%fZ")
                    formatted_release_date = release_date.strftime("%d %B %Y")
                    ea.patch_attribute(attributeId=attribute['attributeId'], value=formatted_release_date)
                if attribute['name'] == "Solved":
                    if challenge['is_owned']:
                        ea.patch_attribute(attributeId=attribute['attributeId'], value="done")
                    else:
                        ea.patch_attribute(attributeId=attribute['attributeId'], value=" ")
                if attribute['name'] == "cssClass":
                    if challenge['is_owned']:
                        ea.patch_attribute(attributeId=attribute['attributeId'], value="done")
                    else:
                        ea.patch_attribute(attributeId=attribute['attributeId'], value="todo")
 
        else: # doesnt already exist, create page
            release_str = challenge['release_date']
            release_date = datetime.strptime(release_str, "%Y-%m-%dT%H:%M:%S.%fZ")
            formatted_release_date = release_date.strftime("%d %B %Y")
            new_note = ea.create_note(
                parentNoteId=catId,
                type="text",
                title=challenge['name'],
                content=" ",
            )
            ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="relation", name="template", value=trilium_challenges_template_id)
            ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="Difficulty", value=challenge['difficulty'])
            ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="Released", value=formatted_release_date)
            if challenge['is_owned']:
                ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="cssClass", value="done")
                ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="Solved", value="done")
            else:
                ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="cssClass", value="todo")
                ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="Solved", value=" ")
 
            print(f"[+] created ID: {challenge['id']}, Name: {challenge['name']}, Difficulty: {challenge['difficulty']}")
 
print("[+] updating folder name                                          ")
ea.patch_note(noteId=trilium_challenges_folder,title="Challenges - "+str(total_completed)+" / "+str(len(challenges)))