#!/usr/bin/env python3
"""Cosecha la metodologia de pentesting que el indexador normal se deja fuera.

    ./cosecha.py            recorre COFRE y escribe datos/cosecha.jsonl
    ./cosecha.py --cuenta   dice que cosecharia, sin escribir

## Por que existe, aparte de indexa.py

`indexa.py` es conservador a proposito: solo mira ficheros cuyo NOMBRE parece
documentacion (readme, notas, guia...) y no baja mas de 3 niveles. Eso funciona
para las herramientas, pero **tira la mejor metodologia que hay en el disco**:

- El `OSCP_Vault` son 72 notas OSCP escritas a mano —`SQL injection.md`,
  `Reverse shell.md`, `LFI a RCE.md`, `msfvenom.md`...— con comandos reales y
  enlaces cruzados. Ninguna entra: el nombre no lleva "notas/guia", y ademas
  viven en `OSCP_APUNTES/OSCP_Vault/02 - Explotacion web/`, a 5 niveles.
- Igual con las cheatsheets y apuntes sueltos de AD, tunneling, privesc.

Medido antes de escribir esto: **0 de 72 notas del Vault estaban en el indice.**

Este cosechador es agresivo donde el otro es prudente: en los perfiles
ofensivos coge TODO .md/.txt que no sea ruido evidente, a la profundidad que
haga falta, y marca de que TEMA es por la carpeta que lo contiene (las del Vault
van numeradas: "01 - Enumeracion", "02 - Explotacion web"...).

No pisa a indexa.py: escribe a un fichero aparte. La union y el deduplicado los
hace reindexa.py.
"""
import argparse
import json
import os
import re
import sys
import unicodedata


def _norm(t):
    t = unicodedata.normalize("NFD", (t or "").lower())
    t = "".join(c for c in t if unicodedata.category(c) != "Mn")
    return re.sub(r"[^a-z0-9 ]", " ", t)

AQUI = os.path.dirname(os.path.abspath(__file__))
RAIZ = os.path.dirname(os.path.dirname(AQUI))          # .../nucleo
COFRE = os.path.expanduser("~/COFRE")
SALIDA = os.path.join(RAIZ, "datos", "cosecha.jsonl")

# Los perfiles ofensivos: aqui vive el saber de pentesting. En estos se cosecha
# a lo ancho. (LINUX se queda para indexa.py, que ya lo hace bien con man pages.)
PERFILES = ["PENTESTERS", "PHISHERS", "REVERSERS", "DEFACERS", "SNEAKERS",
            "HARD-WARERS", "PROTECTORS", "ZAPPERS", "AUTOMATAS"]

# Carpetas de oro: se recorren ENTERAS, sin limite de profundidad, porque su
# valor esta justo en las notas hondas.
ORO = re.compile(r"(OSCP_APUNTES|OSCP_Vault|TELMO_OSCP|RedTeam-Tools|"
                 r"apuntes|cheat|vault|metodolog|notas)", re.I)

# Lo que no se mira nunca: dependencias, control de versiones, binarios.
IGNORA = re.compile(
    r"(^|/)(node_modules|\.git|\.obsidian|vendor|dist|build|__pycache__|"
    r"\.venv|venv|site-packages|target|\.cache|\.github)(/|$)", re.I)

# Ficheros que son plantilla o licencia, no saber.
RUIDO = re.compile(r"(code_of_conduct|contributing|changelog|license|copying|"
                   r"pull_request|issue_template|\.min\.|package-lock)", re.I)

# Fuera de las carpetas de ORO, solo se coge lo que PARECE documentacion por el
# nombre. Es la misma idea que indexa.py, un poco mas ancha (tutorial, writeup,
# walkthrough). Sin esto entra cada .md de cada exploit.
DOC = re.compile(r"(readme|install|usage|apuntes|notas|trucos|howto|tutorial|"
                 r"walkthrough|writeup|write-up|guide|guia|cheat|manual|tips|"
                 r"metodolog|documentation|docs?)", re.I)

# Tope de ficheros por herramienta fuera del ORO. exploitdb tiene 29.925 .md:
# es una base de datos, no unos apuntes, y meterla entera ahoga el oro. Si una
# herramienta pasa de esto, se la trata como repo a granel y solo se coge su
# documentacion de primer nivel (README).
CAP_FICHEROS = 25

MIN_FRAGMENTO = 60
MAX_FRAGMENTO = 1200
MAX_PROFUNDIDAD = 4        # niveles bajo el perfil, salvo en carpetas de ORO

# Tope de tamano por fichero. Una nota de metodologia son unos pocos KB; un .txt
# de 300 MB es una wordlist (SecLists y similares llevan .txt de gigabytes) y
# leerlo entero revienta la maquina —la swap de este equipo es de 976 MiB—. Por
# encima de esto no es saber, es un diccionario, y no se lee.
MAX_BYTES = 512 * 1024


def _profundidad(ruta):
    return ruta.rstrip("/").count("/")


def _tema(ruta):
    """El tema, deducido de la carpeta. Las del Vault van numeradas."""
    for parte in reversed(ruta.split(os.sep)):
        # "02 - Explotacion web" -> "explotacion web"
        m = re.match(r"^\d{2}\s*[-.]\s*(.+)$", parte)
        if m:
            return m.group(1).strip().lower()
    return ""


def _herramienta(ruta, perfil):
    rel = os.path.relpath(ruta, os.path.join(COFRE, perfil)).split(os.sep)
    return rel[0] if len(rel) > 1 else perfil


def _prefija(titulo, seccion):
    """Lleva el titulo de la nota delante de la seccion, salvo que ya lo tenga.

    Sin esto, un fragmento como "## Deteccion time-based\n' AND SLEEP(5)" no
    sabe de que va —¿time-based de que?—. Con "(SQL injection) ..." delante, el
    vector y el cerebro tienen el contexto. Se evita duplicar si el titulo ya
    aparece en la cabecera de la seccion."""
    if titulo and _norm(titulo) not in _norm(seccion[:80]):
        return f"({titulo}) {seccion}"
    return seccion


def trocea(texto):
    """En fragmentos por encabezado, fusionando los cortos y con el titulo.

    Antes se partia por encabezado y se TIRABA toda seccion de menos de 60
    caracteres. Eso perdia justo lo mejor: una seccion "## Deteccion time-based"
    con `' AND SLEEP(5)-- -` son 50 caracteres, y es oro. Ahora una seccion
    corta se PEGA a la siguiente en vez de tirarse, y cada fragmento lleva
    delante el titulo de la nota para no quedar sin contexto.
    """
    texto = re.sub(r"\A---\n.*?\n---\n", "", texto, flags=re.S)
    m = re.search(r"^#\s+(.+)", texto, re.M)
    titulo = m.group(1).strip() if m else ""

    crudas = [s.strip() for s in re.split(r"\n(?=#{1,3}\s)", texto) if s.strip()]

    # fusionar hacia adelante mientras la seccion sea demasiado corta
    fundidas, i = [], 0
    while i < len(crudas):
        sec = crudas[i]
        while len(sec) < MIN_FRAGMENTO and i + 1 < len(crudas):
            i += 1
            sec += "\n\n" + crudas[i]
        fundidas.append(sec)
        i += 1
    # si la ultima quedo corta, se dobla sobre la anterior. Se saca primero y
    # se indexa despues: fundidas.pop() en el lado derecho ya habria acortado
    # la lista y fundidas[-2] se saldria de rango.
    if len(fundidas) >= 2 and len(fundidas[-1]) < MIN_FRAGMENTO:
        cola = fundidas.pop()
        fundidas[-1] += "\n\n" + cola

    for sec in fundidas:
        if len(sec) < MIN_FRAGMENTO:
            continue                          # una nota entera diminuta
        if len(sec) <= MAX_FRAGMENTO:
            yield _prefija(titulo, sec)
        else:
            buf = ""
            for parrafo in re.split(r"\n\s*\n", sec):
                if len(buf) + len(parrafo) > MAX_FRAGMENTO and buf:
                    yield _prefija(titulo, buf.strip())
                    buf = ""
                buf += parrafo + "\n\n"
            if len(buf.strip()) >= MIN_FRAGMENTO:
                yield _prefija(titulo, buf.strip())


def _candidato(ruta, fich, en_oro):
    """Un fichero vale si es texto, no es ruido, no es enorme, y —fuera del
    oro— su nombre parece documentacion."""
    if not fich.lower().endswith((".md", ".txt")):
        return False
    if RUIDO.search(fich):
        return False
    if not en_oro and not DOC.search(fich):
        return False
    try:
        return os.path.getsize(ruta) <= MAX_BYTES
    except OSError:
        return False


def ficheros(perfil):
    """Los ficheros a cosechar de un perfil, con el tope por herramienta.

    Se recorre cada herramienta (subcarpeta de primer nivel) por separado para
    poder contar sus ficheros: si pasa del cap y no es oro, es un repo a granel
    (exploitdb) y solo se coge su README, no sus 30.000 entradas.
    """
    base = os.path.join(COFRE, perfil)
    if not os.path.isdir(base):
        return

    # sueltos en la raiz del perfil (apuntes-ad-pentest.md, trucos_trampas.md)
    for fich in sorted(os.listdir(base)):
        ruta = os.path.join(base, fich)
        if os.path.isfile(ruta) and _candidato(ruta, fich, en_oro=True):
            yield ruta, True

    for herr in sorted(os.listdir(base)):
        raiz_h = os.path.join(base, herr)
        if not os.path.isdir(raiz_h) or herr.startswith("."):
            continue

        candidatos = []
        for raiz, dirs, fichs in os.walk(raiz_h):
            if IGNORA.search(raiz):
                dirs[:] = []
                continue
            en_oro = bool(ORO.search(raiz))
            if not en_oro and _profundidad(raiz) - _profundidad(raiz_h) > MAX_PROFUNDIDAD:
                dirs[:] = []
                continue
            for fich in sorted(fichs):
                ruta = os.path.join(raiz, fich)
                if _candidato(ruta, fich, en_oro):
                    candidatos.append((ruta, en_oro, _profundidad(ruta)))

        del_oro = [c for c in candidatos if c[1]]
        resto = [c for c in candidatos if not c[1]]

        # el oro entra siempre y entero
        for ruta, en_oro, _ in del_oro:
            yield ruta, True

        # el resto, con tope: si es un repo a granel, solo lo mas alto
        if len(resto) > CAP_FICHEROS:
            resto.sort(key=lambda c: c[2])          # los menos hondos primero
            resto = resto[:CAP_FICHEROS]
        for ruta, en_oro, _ in resto:
            yield ruta, False


def construye(solo_cuenta=False):
    salida = []
    for perfil in PERFILES:
        n_doc = n_frag = 0
        for ruta, en_oro in ficheros(perfil):
            try:
                with open(ruta, encoding="utf-8", errors="ignore") as f:
                    texto = f.read()
            except OSError:
                continue
            if len(texto.strip()) < MIN_FRAGMENTO:
                continue
            tema = _tema(ruta)
            herr = _herramienta(ruta, perfil)
            rel = os.path.relpath(ruta, COFRE)
            hubo = False
            for trozo in trocea(texto):
                salida.append({
                    "tipo": "metodologia" if en_oro else "apunte",
                    "herramienta": herr, "perfil": perfil,
                    "fichero": rel, "tema": tema, "texto": trozo,
                })
                n_frag += 1
                hubo = True
            n_doc += hubo
        print(f"  {perfil:14s} {n_doc:5d} docs   {n_frag:6d} fragmentos", flush=True)

    oro = sum(1 for e in salida if e["tipo"] == "metodologia")
    print(f"\n  total: {len(salida)} fragmentos  ({oro} de metodologia)", flush=True)
    if solo_cuenta:
        return 0

    os.makedirs(os.path.dirname(SALIDA), exist_ok=True)
    with open(SALIDA, "w", encoding="utf-8") as f:
        for e in salida:
            f.write(json.dumps(e, ensure_ascii=False) + "\n")
    print(f"  escrito: {SALIDA}", flush=True)
    return 0


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--cuenta", action="store_true")
    a = p.parse_args()
    return construye(solo_cuenta=a.cuenta)


if __name__ == "__main__":
    sys.exit(main())
