145 lines
4.7 KiB
Python
145 lines
4.7 KiB
Python
from typing import Iterable, Literal
|
|
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 FileSystemBlobLoader
|
|
from langchain_community.document_loaders.blob_loaders.schema import Blob, BlobLoader
|
|
from langchain_community.document_loaders import AsyncHtmlLoader
|
|
from langchain.chains import LLMChain
|
|
from langchain.chains import SimpleSequentialChain
|
|
|
|
from dotenv import load_dotenv
|
|
import requests
|
|
import tempfile
|
|
import zipfile
|
|
import io
|
|
import os
|
|
from validators import url
|
|
from time import time
|
|
|
|
load_dotenv()
|
|
|
|
"""
|
|
This application attempts to automatically solve CTF levels for CS492/CS592 Malware Reverse Engineering.
|
|
"""
|
|
|
|
|
|
# def get_rag_chain():
|
|
# return (
|
|
# {"context": retriever | format_docs, "question": RunnablePassthrough()}
|
|
# | prompt
|
|
# | llm
|
|
# | StrOutputParser()
|
|
# )
|
|
|
|
|
|
session = requests.Session()
|
|
|
|
url = "https://cs492.oregonctf.org/"
|
|
|
|
|
|
def get_group(chapter: int):
|
|
if chapter >= 1 and chapter <= 8:
|
|
return "Ch01-08"
|
|
elif chapter >= 11 and chapter <= 13:
|
|
return "Ch11-13"
|
|
elif chapter >= 15 and chapter <= 16:
|
|
return "Ch15-16"
|
|
elif chapter >= 18 and chapter <= 21:
|
|
return "Ch18-21"
|
|
else:
|
|
return False
|
|
|
|
|
|
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 get_file_path(level, group):
|
|
payload = {
|
|
"setname": group,
|
|
}
|
|
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",
|
|
}
|
|
response = session.post(url + "/download", data=payload, headers=headers)
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
with io.BytesIO(response.content) as zip_file:
|
|
with zipfile.ZipFile(zip_file) as z:
|
|
z.extractall(temp_dir)
|
|
extracted_files = os.listdir(temp_dir)
|
|
# print(f"Extracted files: {extracted_files}")
|
|
for file_name in extracted_files:
|
|
if file_name == level:
|
|
temp_file_path = os.path.join(temp_dir, file_name)
|
|
saved_file_path = os.path.join("./temp", file_name)
|
|
os.rename(temp_file_path, saved_file_path)
|
|
return saved_file_path
|
|
return False
|
|
|
|
|
|
def execute(path: str):
|
|
print("path ", path)
|
|
loader = FileSystemBlobLoader(path, show_progress=True)
|
|
prompt = PromptTemplate.from_template(
|
|
"Describe the attached binary file: {file_content}"
|
|
)
|
|
llm = ChatOpenAI(model_name="gpt-4-turbo", temperature=0)
|
|
blobs: Iterable[Blob] = loader.yield_blobs()
|
|
|
|
for b in blobs:
|
|
chain = SimpleSequentialChain(
|
|
prompt=prompt,
|
|
llm=llm,
|
|
)
|
|
# print('b ',b)
|
|
result = chain.run(file_content=b)
|
|
return result
|
|
|
|
|
|
print(
|
|
"This program will attempt to automatically solve CS 492/592 Malware Reverse Engineering CTF levels."
|
|
)
|
|
print("Enter the name of the level (example: Ch03DynA_Ltrace)")
|
|
|
|
|
|
while True:
|
|
try:
|
|
level: str = input("name of binary>> ")
|
|
if level:
|
|
chapter = level[2:4]
|
|
if chapter.strip().isdigit():
|
|
group = get_group(int(chapter))
|
|
if group:
|
|
start_time = time()
|
|
start_session()
|
|
if os.path.isfile(os.path.join("./temp", level)):
|
|
print("found")
|
|
file_path = os.path.join("./temp", level)
|
|
else:
|
|
file_path = get_file_path(level, group)
|
|
if file_path:
|
|
results = execute(file_path)
|
|
print(results)
|
|
else:
|
|
raise Exception("Not a valid file")
|
|
else:
|
|
raise Exception("Not a valid chapter of a possible file")
|
|
else:
|
|
raise Exception("Bad input format (example: Ch03DynA_Ltrace)")
|
|
else:
|
|
break
|
|
except Exception as e:
|
|
print(e)
|
|
break
|