import streamlit as st
from transformers import pipeline

# Загружаем модель для генерации текста
pipe = pipeline("text-generation", model="WhiteRabbitNeo/WhiteRabbitNeo-13B-v1", trust_remote_code=True)

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

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

# Функция для создания эксплойта
def exploit_creation(vulnerability, language):
    st.write(f"Создание эксплойта для уязвимости: {vulnerability} на языке {language}")
    result = pipe(f"Генерация эксплойта для уязвимости {vulnerability} на языке {language}...")
    st.text_area("Генерированный эксплойт", result[0]['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)