import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.io.InputStream;
import java.io.ByteArrayOutputStream;
import java.io.FileWriter;
import java.io.PrintWriter;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * Java agent that applies the CVE-2026-48558 fix to a running SimpleHelp server.
 *
 * SimpleHelp encrypts its com.aem.* classes in simplehelp.jar and decrypts them
 * at load time via its custom SHClassLoader/SecureRunner. A plain replacement
 * class placed in the jar would be "decrypted" again and corrupted, so instead
 * this agent registers a ClassFileTransformer that, after the original
 * OIDCAuthenticator has been decrypted by the classloader, substitutes the
 * fixed OIDCAuthenticator bytes (which add IDTokenVerifier signature/claims
 * verification before oidcSuccess()).
 *
 * Usage:  java -javaagent:fixagent.jar=/path/to/evidence.log -cp lib/simplehelp.jar SimpleHelp
 *
 * The fixed class bytes are bundled as the resource /OIDCAuthenticator.class.
 * The optional agent argument is a path to an evidence log file so that proof
 * survives SimpleHelp's post-startup stdout redirection.
 */
public class FixAgent {
    private static volatile byte[] FIXED_BYTES;
    private static volatile String EVIDENCE_LOG;
    private static final AtomicBoolean REPLACED = new AtomicBoolean(false);

    public static void premain(String agentArgs, Instrumentation inst) {
        EVIDENCE_LOG = agentArgs;
        if (EVIDENCE_LOG != null && !EVIDENCE_LOG.isEmpty()) {
            System.setProperty("shfix.evidence", EVIDENCE_LOG);
        }
        FIXED_BYTES = loadResource("/OIDCAuthenticator.class");
        if (FIXED_BYTES == null) {
            evidence("FATAL: fixed OIDCAuthenticator.class resource not found");
            return;
        }
        evidence("FixAgent premain: fixed OIDCAuthenticator resource loaded (" + FIXED_BYTES.length + " bytes)");
        inst.addTransformer(new ClassFileTransformer() {
            @Override
            public byte[] transform(Module module, ClassLoader loader, String className,
                                    Class<?> classBeingRedefined,
                                    ProtectionDomain protectionDomain,
                                    byte[] classfileBuffer) {
                if ("com/aem/shelp/proxy/authentication/oidc/OIDCAuthenticator".equals(className)) {
                    if (REPLACED.compareAndSet(false, true)) {
                        evidence("AGENT_APPLIED: replaced OIDCAuthenticator with CVE-2026-48558 fixed version (loader=" + loader + ")");
                    }
                    byte[] copy = new byte[FIXED_BYTES.length];
                    System.arraycopy(FIXED_BYTES, 0, copy, 0, FIXED_BYTES.length);
                    return copy;
                }
                return null;
            }
        }, false);
    }

    static void evidence(String msg) {
        String line = "[FixAgent " + System.currentTimeMillis() + "] " + msg;
        System.out.println(line);
        if (EVIDENCE_LOG == null || EVIDENCE_LOG.isEmpty()) return;
        try (PrintWriter pw = new PrintWriter(new FileWriter(EVIDENCE_LOG, true))) {
            pw.println(line);
        } catch (Exception e) {
            System.err.println("[FixAgent] failed to write evidence: " + e);
        }
    }

    private static byte[] loadResource(String name) {
        try (InputStream in = FixAgent.class.getResourceAsStream(name)) {
            if (in == null) return null;
            ByteArrayOutputStream bout = new ByteArrayOutputStream();
            byte[] buf = new byte[8192];
            int n;
            while ((n = in.read(buf)) > 0) bout.write(buf, 0, n);
            return bout.toByteArray();
        } catch (Exception e) {
            System.err.println("[FixAgent] Failed to load resource " + name + ": " + e);
            return null;
        }
    }
}