#core/CLI.py

import subprocess
import typer

app = typer.Typer()

try:
    @app.command()
    def up():
        subprocess.run(
            ["docker", "compose", "up", "--build"],
            check=True
        )
    
    @app.command()
    def down():
        subprocess.run(
            ["docker", "compose", "down"],
            check=True
        )
        
    @app.command()
    def logs():
        subprocess.run(
            ["docker", "compose", "logs", "-f"],
            check=True
        )
        
    @app.command()
    def reanudar():
        subprocess.run(
            ["docker", "compose", "start"],
            check=True
        )
        
    @app.command()
    def detener():
        subprocess.run(
            ["docker", "compose", "stop"],
            check=True
        )
except subprocess.CalledProcessError:
    pass


except FileNotFoundError:
    typer.echo("Docker no está instalado o no está en el PATH.")
    

if __name__ == "__main__":
    app()


#___________________________________
#core/main.py
from fastapi import FastAPI
from app.start import Inicio
import traceback
from contextlib import (
    asynccontextmanager
)


@asynccontextmanager
async def lifespan(app: FastAPI):
    startup_inicio = Inicio()
    
    print("iniciando app")
    
    try:
        await startup_inicio.startup()
        yield
        
    except Exception as e:
        print(e)
        traceback.print_exc()
        
    finally:
        await startup_inicio.shutdown()


app = FastAPI(
    lifespan=lifespan
)



@app.get("/")
async def running():
    return {"status": "ok"}
    
    

#___________________________________
#app/start.py
from app.socket_run import server_socket
#from app.socket_run import server_init #funcion
from temporal.temporal_register import workerflow_register # funcion
from temporal.start_worker import Workflow_start #clase
from app.base_sql import SQL 
import asyncio
import traceback


class Inicio:
    def __init__(self):
        self.workflow_start = Workflow_start()
        self.SQL = SQL
        self.worker_ready = asyncio.Event()
        self.tasks = []
        self.queue = asyncio.Queue()

    async def startup(self):
        print("iniciando conexiones")

        await self.SQL.create_tables()

        #self.tasks.append(asyncio.create_task(server_init(self.queue)))
        sele.tasks.append(asyncio.create_task(server_socket(self.queue)))
        
        self.tasks.append(asyncio.create_task(workerflow_register(self.worker_ready)))

        try:
            await asyncio.wait_for(
                self.worker_ready.wait(),
                timeout=30
            )        
            
        except asyncio.TimeoutError as e:
            print(f"error: {e}")
            traceback.print_exc()
            raise RuntimeError("Worker no disponible")

        
        self.tasks.append(asyncio.create_task(self.workflow_start.tasks_workflow(self.queue)))
        
        
    async def shutdown(self):
        await self.SQL.close()

        for task in self.tasks:
            if not task.done():
                task.cancel()

        await asyncio.gather(
            *self.tasks,
            return_exceptions=True
        )

#___________________________________
#core/base_module.py
from abc import ABC, abstractmethod

class BaseModule(ABC):
    @abstractmethod
    async def run(self, *args, **kwargs):
        ...


#___________________________________
#app/socket_run.py
import asyncio
import traceback
from functools import partial

#async def server_init(queue):
    
   # task_server = asyncio.create_task(server_socket(queue))
    
    #try:
       # await task_server
        
    #finally:
        #if not task_server.done():
          #  task_server.cancel()
    
       # await asyncio.gather(
          #  task_server,
         #   return_exceptions=True
       # )  



async def server_socket(queue):
    server = None
    
    try:
        server = await asyncio.start_server(
            partial(handle_client, queue),
            "0.0.0.0",
            9000
        )

        print("corriendo servidor")

        async with server:
            await server.serve_forever()

    except asyncio.CancelledError:
        print("servidor cancelado")
        raise
        
    finally:
        print("cerrando servidor")
        if server is not None:
            server.close()
            await server.wait_closed()
            
              

async def handle_client(queue, reader, writer):
    
    host = writer.get_extra_info("peername")

    print(f"host conectado: {host}")

    try:
        data = await asyncio.wait_for(
            reader.read(1024),
            timeout=30
        )

        if not data:
            print("conexión sin datos")
            return
            
        print(data.decode(errors="replace"))

        await queue.put(data.decode())
        
        writer.write(b"ok")

        try:
            await writer.drain()
        except (ConnectionResetError, BrokenPipeError):
            print(f"host caido: {host}")
            return

    except asyncio.CancelledError:
        print(f"cliente cancelado: {host}")
        raise
        
        
    except asyncio.TimeoutError:
        print("no se envio datos...")
        print(f"cerrando host: {host}")
        return

    finally:
        print("cerrando host")

        writer.close()

        try:
           await writer.wait_closed() 
        except Exception as e:
            print(e)
            traceback.print_exc()
            

#___________________________________
#temporal/temporal_register.py
from temporalio.client import Client
from temporalio.worker import Worker
from temporalio.service import RPCError
import asyncio
import traceback

from .workflows import MainWorkflow
from .activity_worker import sample_activity

async def workerflow_register(worker_ready):
    try:
        client = await Client.connect(
            "temporal:7233"
        )
        
    except RPCError as e:
        print(f"error al conectar temporal: {e}")
        raise
        
    
    worker = Worker(
        client,
        task_queue = "framework",
        workflows = [MainWorkflow],
        activities = [sample_activity]
    )

    health = asyncio.create_task(state_health(client))
    
    try:
        await worker.run()

        worker_ready.set()
        
    except Exception as e:
        print(f"error: {e}")
        traceback.print_exc()
        await asyncio.gather(health, return_exceptions=True)
        raise

    finally:
        health.cancel()
        await asyncio.gather(health, return_exceptions=True)


async def state_health(client):
    while True:
        try:
            await client.workflow_service.get_system_info()
            
        except asyncio.CancelledError:
            raise
            
        except Exception as e:
            print(f"temporal caido: {e}")
            traceback.print_exc()

        finally:
            await asyncio.sleep(30)
            


#___________________________________
#temporal/start_worker.py
from temporalio.client import Client
from .workflows import MainWorkflow
import uuid
from temporalio.service import RPCError

class Workflow_start:
    def __init__(self):
        self.client = None


    async def initialize_temporal(self):
        try:
            self.client = await Client.connect(
                "temporal:7233"
            )
        except RPCError as e:
            print(f"error al conectar temporal: {e}")
            raise
            
        
        
    async def tasks_workflow(self, queue):
        result = None

        while True:
            
            dato = await queue.get()
            
            try:
                if self.client is None:
                    await self.initialize_temporal()

                result = await self.client.execute_workflow(
                    MainWorkflow.run,
                    dato,
                    id = f"mainworkflow-{uuid.uuid4()}",
                    task_queue = "framework"
                )
            
            except RPCError as e:
                print(f"error al ejecutar .execute_workflow: {e}")

            except asyncio.CancelledError:
                raise

            except Exception as e:
                print(f"error: {e}")
                raise

            finally:
                queue.task_done()

        print(result)

#___________________________________
#temporal/workflows.py
from .activity_worker import sample_activity
from datetime import timedelta
from temporalio import workflow

@workflow.defn
class MainWorkflow:
    @workflow.run
    async def run(self, dato: str):
        try:
            
            result = await workflow.execute_activity(
                sample_activity,
                dato,
                start_to_close_timeout=timedelta(minutes=1)
            )

            return result
            
        except Exception as e:
            print(f"error en .execute_activity: {e}")
            raise
            
#___________________________________
#temporal/activity_worker.py
from temporalio import activity
from app.graph_main import grafo
import logging

logger = logging.getLogger(__name__)

@activity.defn
async def sample_activity(dato):
    
    try:
        return await grafo(dato)
        
    except Exception as e:
        logger.exception(e)
        raise
        
    



#___________________________________
#app/graph_main.py
from app.graph_start import graph_init
import logging

logger = logging.getLogger(__name__)


async def grafo(entrada):
    
    graph = await graph_init()

    try:
        
        result = await graph.ainvoke({"data": entrada})
        
    except Exception as e:
        logger.exception(f"error: {e}")
        raise

    return result



#___________________________________
#app/graph_start.py
from langgraph.graph import StateGraph
from herramientas.gobuster import h_gobuster 
from herramientas.dirsearch import analizer_dirsearch
from herramientas.ffuf import analizer_ffuf 

class state(dict):
    pass
    

async def graph_init():
    Graph = StateGraph(state)
        
    Graph.add_node("gobuster", gobuster)
    Graph.add_node("dirsearch", dirsearch)
    Graph.add_node("ffuf", ffuf)

    Graph.add_edge("gobuster", "dirsearch")
    Graph.add_edge("dirsearch", "ffuf")

    Graph.set_entry_point("gobuster")

    Graph.set_finish_point("ffuf")

    return Graph.compile()




async def gobuster(states):
    print("iniciando herramienta gobuster")

    resultado_gobuster = await h_gobuster(state)

    state["resultado_gobuster"] = resultado_gobuster

    return state



async def dirsearch(state):
    print("iniciando herramienta dirsearch")

    resultado_dirsearch = await analizer_dirsearch(states) 

    state["resultado_dirsearch"] = resultado_dirsearch

    return state

    

async def ffuf(state):
    print("iniciando herramienta ffuf")

    resultado_ffuf = await analizer_ffuf(states) #analizer()

    state["resultado_ffuf"] = resultado_ffuf

    return state

        

#___________________________________
#herramientas/gobuster.py
from app.base_sql import SQL, Gobuster
import asyncio
import re
import logging

logger = logging.getLogger(__name__)

async def h_gobuster(dict_url: str):
    url = dict_url.get("data", "")
    resultado = None

    if not url:
        print("falta url de entrada")
        return []

    comando = ["gobuster", "dir", "-u", f"{url}", "-w", "wordlist.txt"]

    try:
        resultado = await asyncio.create_subprocess_exec(
            *comando,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )

        stdout, stderr = await asyncio.wait_for(
            resultado.communicate(),
            timeout=60
        )

        if resultado.returncode != 0:
            print(f"error: {stderr.decode()}")
            return []

        
    except asyncio.CancelledError:
        logger.exception("proceso cancelado")
        raise

    except FileNotFoundError:
        logger.exception("herramienta no encontrada")
        return []

    except asyncio.TimeoutError:
        logger.exception("tiempo agotado")
        return []

    except Exception as e:
        logger.exception(f"error: {e}")
        return []

    output = stdout.decode(errors="ignore")

    coincidencia = re.findall(r"^(\/\S+)", output, re.MULTILINE)

    #data = {url: {"directorios": coincidencia}}

    try:
        await SQL.save_data(Gobuster(url={url: {"directorios": coincidencia}}))
        
    except Exception as e:
        logger.exception(f"error: {e}")
        
    
    
    return coincidencia


#___________________________________
#herramientas/dirsearch.py
from app.base_sql import SQL, Dirsearch
import asyncio
import re
import logging

logger = logging.getLogger(__name__)

sem = asyncio.Semaphore(5)

async def analizer_dirsearch(dict_url):
    tasks = []

    url = dict_url.get("data", "")
    data = dict_url.get("resultado_gobuster", {})
    data_dict = {}

    if not data:
        print("sin directorios de gobuster")
        return {}

    for d in data:
        tasks.append(asyncio.create_task(h_dirsearch(url, d)))


    resultado = await asyncio.gather(*tasks)

    for directorio, coincidencia in resultado:
        data_dict[directorio] = coincidencia

    try:
        await SQL.save_data(Dirsearch(url={url: data_dict}))
        
    except Exception as e:
        logger.exception(f"error: {e}")
    
    return data_dict
        


async def h_dirsearch(url: str, d: str):
    async with sem:
        comando = ["dirsearch", "-u", f"{url.rstrip('/')}/{d.strip('/')}/"]

        try:
            resultado = await asyncio.create_subprocess_exec(
                *comando,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
                
            stdout, stderr = await asyncio.wait_for(
                resultado.communicate(),
                timeout=60
            )

            if resultado.returncode != 0:
                print(f"error: {stderr.decode()}")
                return d, []
            
        except asyncio.TimeoutError:
            logger.warning("tiempo agotado")
            return d, []

        except asyncio.CancelledError:
            logger.info("proceso cancelado")
            raise

        except FileNotFoundError:
            logger.error("herramientas no encontrada")
            return d, []

        except Exception as e:
            logger.exception(f"error: {e}")
            return d, []

        output = stdout.decode(errors="ignore")

        coincidencia = re.findall(r"^\/\S+", output, re.MULTILINE)

        return d, coincidencia      


#___________________________________
#herramientas/ffuf.py
from app.base_sql import SQL, Ffuf
import asyncio
import re
import logging

logger = logging.getLogger(__name__)

sem = asyncio.Semaphore(5)

async def analizer_ffuf(dict_url):
    tasks = []

    url = dict_url.get("data", "")
    data = dict_url.get("resultado_dirsearch", {})
    info = {}

    if not data:
        logger.warning("datos de dirsearch vacios")
        return {}

    if not url:
        logger.error("url vacia")
        return {}

    for clave, valor in data.items():
        if not isinstance(valor, list):
            logger.warning(f"{clave} no contiene un lista")
            continue
            
        for lista in valor:
            tasks.append(asyncio.create_task(h_ffuf(url, clave, lista)))

    if not tasks:
        return {}
        
    resultado = await asyncio.gather(*tasks, return_exceptions=True
)


    for item in resultado:
        if isinstance(item, Exception):
            logger.error(f"Tarea falló: {item}")
            continue

        clave, listas, resultados = item
        
        info.setdefault(clave, {})[listas] = resultados
        

    try:
        await SQL.save_data(
            Ffuf(
                url={url: info}
            )
        )
            
    except Exception as e:
        logger.exception(f"error: {e}")
        
    return info



async def h_ffuf(url: str, clave: str, lista: str):
    info_list = []
    
    async with sem:
        comando = ["ffuf", "-u", f"{url.rstrip('/')}/{lista.strip('/')}/FUZZ", "-w", "wordlist.txt"]

        try:
            proceso = await asyncio.create_subprocess_exec(
                *comando,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
                    
            stdout, stderr = await asyncio.wait_for(
                proceso.communicate(),
                timeout=60
            )

            if proceso.returncode !=0:
                logger.error(stderr.decode(errors="ignore"))
                return clave, lista, []
                    
                
        except asyncio.TimeoutError:
            logger.warning("tiempo agotado")
            return clave, lista, []

        except asyncio.CancelledError:
            logger.info("proceso cancelado")
            raise

        except FileNotFoundError:
            logger.error("no se encontró la herramienta")
            return clave, lista, []

        except Exception as e:
            logger.exception(f"error: {e}")
            return clave, lista, []

        output = stdout.decode(errors="ignore")

        for match in re.finditer(r"^(\S+)\s+\[Status:\s*(\d+).*?\]$", output, re.MULTILINE):
            line = match.group(0)
            status = int(match.group(2))

            if status in [404, 400]:
                continue
                    
            info_list.append(line)

        return clave, lista, info_list


                
#___________________________________
#app/base_sql.py
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from sqlalchemy import JSON
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker

class Base(DeclarativeBase):
    pass

#tabla
class Gobuster(Base):
    __tablename__ = "gobuster"

    id: Mapped[int] = mapped_column(primary_key=True)

    url: Mapped[dict] = mapped_column(JSON)


#tabla
class Dirsearch(Base):
    __tablename__ = "dirsearch"

    id: Mapped[int] = mapped_column(primary_key=True)

    url: Mapped[dict] = mapped_column(JSON)


#tabla
class Ffuf(Base):
    __tablename__ = "ffuf"

    id: Mapped[int] = mapped_column(primary_key=True)

    url: Mapped[dict] = mapped_column(JSON)
            

class Base_sql:
    def __init__(self):
        self.DATABASE_URL = (
            "postgresql+asyncpg://postgres:postgres@postgres:5432/framework"
        )
        self.engine = create_async_engine(
            self.DATABASE_URL,
            echo=True
        )
        self.SessionLocal = async_sessionmaker(
            self.engine,
            expire_on_commit=False
        )
    

    async def create_tables(self):
        async with self.engine.begin() as conn:
            await conn.run_sync(
                Base.metadata.create_all
            )

    
    async def close(self):
        await self.engine.dispose()
        

    
    async def save_data(self, obj):
        async with self.SessionLocal as sesion:
            try:
                sesion.add(obj)
                await sesion.commit()
                await sesion.refresh(obj)
            except Exception:
                await sesion.rollback()
                raise
            
                
SQL = Base_sql()