import gradio as gr
from huggingface_hub import InferenceClient

# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("WhiteRabbitNeo/Trinity-33B-v1.0")
model = AutoModelForCausalLM.from_pretrained("WhiteRabbitNeo/Trinity-33B-v1.0")


# Создание клиента для доступа к модели
client = InferenceClient(model_id,)

# Основная функция для разведки
def reconnaissance(domain):
    st.write(f"Начинаем разведку для домена: {domain}")
    result = client.predict(f"Разведка информации о домене {domain}...")
    st.text_area("Результаты разведки", result <sup> </sup>['generated_text'])

# Функция для поиска уязвимостей
def vulnerability_search(service_info):
    st.write(f"Ищем уязвимости для сервиса: {service_info}")
    result = client.predict(f"Поиск уязвимостей для сервиса {service_info}...")
    st.text_area("Результаты поиска уязвимостей", result <sup> </sup>['generated_text'])

# Функция для создания эксплойта
def exploit_creation(vulnerability, language):
    st.write(f"Создание эксплойта для уязвимости: {vulnerability} на языке {language}")
    result = client.predict(f"Генерация эксплойта для уязвимости {vulnerability} на языке {language}...")
    st.text_area("Генерированный эксплойт", result <sup> </sup>['generated_text'])

# Интерфейс Streamlit
st.title("Инструмент для Red Team операций с WhiteRabbitNeo")
st.markdown("Используйте этот инструмент для разведки, поиска уязвимостей и создания эксплойтов.")

# Выбор задачи
task = st.sidebar.selectbox("Выберите задачу:", ["Разведка", "Поиск уязвимостей", "Создание эксплойта"])

if task == "Разведка":
    domain = st.text_input("Введите домен для разведки:")
    if st.button("Начать разведку"):
        reconnaissance(domain)

elif task == "Поиск уязвимостей":
    service_info = st.text_input("Введите информацию о сервисе (например, Apache 2.4.49):")
    if st.button("Искать уязвимости"):
        vulnerability_search(service_info)

elif task == "Создание эксплойта":
    vulnerability = st.text_input("Опишите уязвимость:")
    language = st.selectbox("Выберите язык для эксплойта:", ["Python", "Bash", "PowerShell"])
    if st.button("Создать эксплойт"):
        exploit_creation(vulnerability, language)

# Запуск приложения
if __name__ == "__main__":
    demo.launch()