from langchain_community.document_loaders.generic import GenericLoader
from langchain_community.document_loaders.parsers import LanguageParser
from langchain_openai import ChatOpenAI
from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_community.document_loaders import AsyncHtmlLoader
from dotenv import load_dotenv
import requests
from validators import url
from time import time

"""
This application attempts to automatically solve CTF levels for CS492/CS592 Malware Reverse Engineering.
"""
load_dotenv()
promt = "You are a malware reverse engineer"
def get_rag_chain():
    return (
        {"context": retriever | format_docs, "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )

session = requests.Session()

llm = ChatOpenAI(model_name="gpt-4-turbo", temperature=0)
url = "https://cs492.oregonctf.org/"

def start_session():
    payload = {
        'username': 'demo0',
        'passwd': 'malware'
    }

    headers = {
        'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7',
        'Content-Type': 'application/x-www-form-urlencoded'
    }
    session.post(url, data=payload, headers=headers)

def download(setname):
    payload = {
        'setname': setname,
    }
    headers = {
        'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7',
        'Content-Type': 'application/x-www-form-urlencoded'
    }
    p = session.post(url+'/download', data=payload, headers=headers)
    print(p.headers.get('Content-Type'))
    zip = p.raw
    print(zip)



while True:
    try:
        start_session()
        print(session.cookies)
        download('Ch01-08')
        line: str = input("llm>> ")
        if line:
            start_time = time()
         
            result: str = get_rag_chain.invoke(line)
            end_time = time()
            elapsed_time = round(end_time - start_time, 2)
            print("\n", result, "\n\nElapsed time: ", elapsed_time, " seconds")
        else:
            break
    except Exception as e:
        print(e)
        break