chat.py

from autogdb import *
from autogdb.llm import LLMConfig
import json
import os

from rich import print
from rich.progress import Progress
import readline


import warnings


warnings.filterwarnings("ignore", message="You are trying to use a chat model.*")
from rich import print

CACHE_FILE_PATH = '.server_cache.auotogpt.json'


lo = Logger()

def banner():
    banner = """\
    \n\n
           _____                _____)  ______    ______   
          (, /  |             /        (, /    ) (, /    ) 
            /---|     _/_ ___/   ___     /    /    /---(   
         ) /    |_(_(_(__(_)/     / )  _/___ /_ ) / ____)  
        (_/                (____ /   (_/___ /  (_/ ("""
    
    author = """\n\
        Enable GPT in your reversing job.
        [bold red]>[/bold red] Author [bold blue]Retr0Reg[/bold blue], ChatWithBinary Team.       
    """
    print(banner,end='')
    print(author,end='')
    print('\n')

def load_llm_config() -> LLMConfig:
    config = LLMConfig.from_env()
    lo.success("Loaded LLM config from environment.")

    if config.api_key is None:
        lo.fail("API key not found in environment.")
        config.api_key = input("Please enter your LLM API key: ").strip()
        if config.provider == "openai":
            config.api_base = (
                input("Please enter the OpenAI API base URL: ").strip()
                or config.api_base
            )

        with open('api_key.py', 'w') as file:
            file.write(f'OPENAI_API_KEY = "{config.api_key}"\n')
            file.write(f'OPENAI_API_BASE = "{config.api_base}"\n')

        lo.success("api_key.py file created with your API key and base URL.")

    return config

def console_input(input_str: str) -> str:
    print(input_str, end='')
    return input()

def get_server_info():
    if os.path.exists(CACHE_FILE_PATH):
        with open(CACHE_FILE_PATH, 'r') as cache_file:
            try:
                server_info = json.load(cache_file)
                addr = server_info['ip']
                port = server_info['port']
                
                if (not addr) or (not port):
                    raise KeyError(f"Server address and port saved is empty, please save it again by deleting:{CACHE_FILE_PATH}")
                
                return addr,port
                
            except json.JSONDecodeError:
                lo.fail("Cache file is corrupted. Please enter server details again.")

    # Cache file doesn't exist or is corrupted, ask the user for info
    server_ip = console_input("[bold light_steel_blue1][?] Please enter your server IP:[/bold light_steel_blue1] ").strip()
    server_port = console_input("[bold light_steel_blue1][?] Please enter your server port:[/bold light_steel_blue1] ").strip()
    try:
        with open(CACHE_FILE_PATH, 'w') as cache_file:
            json.dump({'ip': server_ip, 'port': server_port}, cache_file)
        lo.success(f"Server address and port saved, change it by deleting {CACHE_FILE_PATH}.")
    except Exception as e:
        lo.fail("Save address and port failed, please check for your privilege and etc.")

    return server_ip, server_port

def setup():
    llm_config = load_llm_config()
    ip, port = get_server_info()
    pwnagent = PwnAgent(
        llm_config.api_key,
        llm_config.api_base,
        AutoGDB(server=ip, port=port).tool(),
        llm_config=llm_config,
    )
    chatagent = ChatAgent(
        llm_config.api_key,
        llm_config.api_base,
        pwnagent,
        llm_config=llm_config,
    )
    return chatagent

def cli():
    chatagent = setup()
    while True:
        text_query = console_input(f"\n  [bold light_steel_blue1] Talk to [/bold light_steel_blue1][bold plum2]GDBAgent[/bold plum2]>>> ")
        print(f"  [bold medium_purple1]:snowboarder: GDBAgent[/bold medium_purple1]: ", end='')
        chatagent.chat_and_assign(text_query)

if __name__ == "__main__":
    banner()
    cli()