Update app.py

app.py

@@ -1,72 +1,46 @@
 import streamlit as st
-import os
-
-# Инициализация истории сессий
-if "recon_history" not in st.session_state:
-    st.session_state.recon_history = []
-
-# Заголовок приложения
-st.title("WhiteRabbitNeo: Red Team Interface")
-
-# Раздел для ввода IP-адреса или домена для разведки
-st.header("Разведка цели")
-target = st.text_input("Введите IP-адрес или домен цели:")
-
-# Пример данных разведки
-recon_data = f"""
-- Цель: {target}
-- IP-адрес: 192.168.1.1
-- Открытые порты: 22, 80, 443
-- Сервисы: SSH, HTTP, HTTPS
-"""
-
-# Кнопка для запуска разведки
-if st.button("Запустить разведку"):
-    if target:
-        st.success("Разведка выполнена успешно!")
-        st.subheader("Результаты разведки")
-        st.code(recon_data, language="text")
-    else:
-        st.error("Пожалуйста, введите IP-адрес или домен.")
-
-# Сохранение результатов разведки в файл
-if st.button("Сохранить результаты разведки"):
-    with open("recon_results.txt", "w") as f:
-        f.write(recon_data)
-    st.success("Результаты разведки сохранены в файл recon_results.txt")
-    st.download_button(label="Скачать файл результатов разведки", data=recon_data, file_name="recon_results.txt", mime="text/plain")
-
-# Сохранение результатов разведки в историю сессии
-if st.button("Сохранить результаты в сессию"):
-    st.session_state.recon_history.append(recon_data)
-    st.success("Результаты разведки добавлены в историю сессии.")
-
-# Отображение истории сессий
-st.subheader("История разведки")
-for idx, data in enumerate(st.session_state.recon_history):
-    st.write(f"Сессия {idx+1}:")
-    st.code(data, language="text")
-
-# Пример кода эксплойта
-st.header("Сгенерированный код эксплойта")
-exploit_code = """
-import socket
-
-def exploit(target_ip, target_port):
-    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-    s.connect((target_ip, target_port))
-    s.send(b"EXPLOIT_PAYLOAD")
-    response = s.recv(1024)
-    s.close()
-    return response
-"""
-
-# Отображение сгенерированного кода эксплойта в формате Python
-st.code(exploit_code, language="python")
-
-# Кнопка для сохранения кода эксплойта в файл
-if st.button("Сохранить код эксплойта"):
-    with open("exploit_code.py", "w") as f:
-        f.write(exploit_code)
-    st.success("Код эксплойта сохранен в файл exploit_code.py")
-    st.download_button(label="Скачать код эксплойта", data=exploit_code, file_name="exploit_code.py", mime="text/x-python")
+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)