/*_############################################################################
  _##
  _##  SNMP4J-Agent 3 - AgentStateDeserializationFilterTest.java
  _##
  _##  Copyright (C) 2005-2026  Frank Fock (SNMP4J.org)
  _##
  _##  Licensed under the Apache License, Version 2.0 (the "License");
  _##  you may not use this file except in compliance with the License.
  _##  You may obtain a copy of the License at
  _##
  _##      http://www.apache.org/licenses/LICENSE-2.0
  _##
  _##  Unless required by applicable law or agreed to in writing, software
  _##  distributed under the License is distributed on an "AS IS" BASIS,
  _##  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  _##  See the License for the specific language governing permissions and
  _##  limitations under the License.
  _##
  _##########################################################################*/

package org.snmp4j.agent.io;

import org.junit.Test;
import org.snmp4j.smi.*;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.InvalidClassException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.net.InetAddress;
import java.util.Collections;
import java.util.HashMap;

import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.fail;

/**
 * Tests the strict {@link AgentStateDeserializationFilter} applied by {@link DefaultMOInput}: legitimate
 * agent state classes round-trip, while disallowed classes, disallowed nested objects, disallowed
 * arrays and overly deep graphs are rejected.
 *
 * @author Frank Fock
 */
public class AgentStateDeserializationFilterTest {

    private static byte[] serialize(Object o) throws Exception {
        ByteArrayOutputStream bos = new ByteArrayOutputStream();
        try (ObjectOutputStream oos = new ObjectOutputStream(bos)) {
            oos.writeObject(o);
        }
        return bos.toByteArray();
    }

    /**
     * Wraps the serialized bytes in an {@link ObjectInputStream} that is protected by the default strict
     * filter (the same way {@link DefaultMOInput} installs it).
     */
    private static ObjectInputStream filteredStream(byte[] bytes) throws Exception {
        ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(bytes));
        // DefaultMOInput installs the strict filter as a side effect of construction.
        new DefaultMOInput(ois);
        return ois;
    }

    @Test
    public void testRoundTripAllowedTypesThroughDefaultMOInput() throws Exception {
        Context context = new Context(new OctetString("ctx"));
        MOInfo moInfo = new MOInfo(new OID("1.3.6.1.2.1.1.1"), "1.0");
        Sequence sequence = new Sequence(42);
        Variable[] values = new Variable[]{
                new Integer32(7),
                new OctetString(new byte[]{1, 2, 3, 4}),
                new OID("1.3.6.1.4.1.4976"),
                new Null(),
                new Counter64(123456789L),
                new TimeTicks(98765L),
                new Gauge32(55L),
                new Counter32(66L),
                new Opaque(new byte[]{9, 8, 7}),
                new IpAddress(InetAddress.getByName("127.0.0.1")),
                new IpAddress(InetAddress.getByName("::1")),
                new UdpAddress(InetAddress.getByName("10.0.0.1"), 161),
                new TcpAddress(InetAddress.getByName("10.0.0.2"), 162)
        };
        IndexedVariables indexed = new IndexedVariables(new OID("1.2"), values);

        // Write all model objects into a single stream, mimicking DefaultMOOutput.
        ByteArrayOutputStream bos = new ByteArrayOutputStream();
        try (ObjectOutputStream oos = new ObjectOutputStream(bos)) {
            oos.writeObject(context);
            oos.writeObject(moInfo);
            oos.writeObject(sequence);
            oos.writeObject(indexed);
        }

        ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(bos.toByteArray()));
        DefaultMOInput in = new DefaultMOInput(ois);

        Context readContext = in.readContext();
        assertEquals(new OctetString("ctx"), readContext.getContext());

        MOInfo readMOInfo = in.readManagedObject();
        assertEquals(new OID("1.3.6.1.2.1.1.1"), readMOInfo.getOID());

        Sequence readSequence = in.readSequence();
        assertEquals(42, readSequence.getSize());

        IndexedVariables readIndexed = in.readIndexedVariables();
        assertNotNull(readIndexed);
        assertEquals(new OID("1.2"), readIndexed.getIndex());
        assertEquals(values.length, readIndexed.getValues().length);
        assertEquals(new Integer32(7), readIndexed.getValues()[0]);
        assertArrayEquals(new byte[]{1, 2, 3, 4}, ((OctetString) readIndexed.getValues()[1]).getValue());
        assertEquals(new OID("1.3.6.1.4.1.4976"), readIndexed.getValues()[2]);
        assertEquals(new Counter64(123456789L), readIndexed.getValues()[4]);
        assertEquals(new IpAddress(InetAddress.getByName("127.0.0.1")), readIndexed.getValues()[9]);
        assertEquals(new IpAddress(InetAddress.getByName("::1")), readIndexed.getValues()[10]);
        assertEquals(new UdpAddress(InetAddress.getByName("10.0.0.1"), 161), readIndexed.getValues()[11]);
    }

    @Test
    public void testStandaloneVariableRoundTrip() throws Exception {
        OctetString os = new OctetString("hello world");
        DefaultMOInput in = new DefaultMOInput(new ObjectInputStream(new ByteArrayInputStream(serialize(os))));
        Variable v = in.readVariable();
        assertEquals(os, v);
    }

    @Test
    public void testDisallowedTopLevelClassRejected() throws Exception {
        HashMap<String, String> map = new HashMap<>();
        map.put("a", "b");
        byte[] bytes = serialize(map);
        ObjectInputStream ois = filteredStream(bytes);
        try {
            ois.readObject();
            fail("Expected InvalidClassException for disallowed HashMap");
        }
        catch (InvalidClassException expected) {
            // expected: filter rejected java.util.HashMap
        }
    }

    @Test
    public void testNestedDisallowedClassRejected() throws Exception {
        Object[] payload = new Object[]{new Integer32(1), new HashMap<String, String>()};
        byte[] bytes = serialize(payload);
        ObjectInputStream ois = filteredStream(bytes);
        try {
            ois.readObject();
            fail("Expected InvalidClassException for Object[] containing a HashMap");
        }
        catch (InvalidClassException expected) {
            // expected: both Object[] and the nested HashMap are disallowed
        }
    }

    @Test
    public void testPrimitiveArraysAllowedButObjectArraysRejected() throws Exception {
        // primitive arrays are allowed
        ObjectInputStream okBytes = filteredStream(serialize(new byte[]{1, 2, 3}));
        assertArrayEquals(new byte[]{1, 2, 3}, (byte[]) okBytes.readObject());
        ObjectInputStream okInts = filteredStream(serialize(new int[]{4, 5, 6}));
        assertArrayEquals(new int[]{4, 5, 6}, (int[]) okInts.readObject());

        // String[] is not allowed (String is allowed, but a String array is not part of the model)
        ObjectInputStream stringArray = filteredStream(serialize(new String[]{"x", "y"}));
        try {
            stringArray.readObject();
            fail("Expected InvalidClassException for String[]");
        }
        catch (InvalidClassException expected) {
            // expected
        }
    }

    @Test
    public void testDepthLimitRejected() throws Exception {
        // Build a deeply nested array-of-arrays of byte that exceeds MAX_DEPTH.
        Object nested = new byte[]{0};
        for (int i = 0; i < AgentStateDeserializationFilter.MAX_DEPTH + 2; i++) {
            Object[] wrapper = new Object[]{nested};
            nested = wrapper;
        }
        byte[] bytes = serialize(nested);
        ObjectInputStream ois = filteredStream(bytes);
        try {
            ois.readObject();
            fail("Expected InvalidClassException due to exceeded depth limit");
        }
        catch (InvalidClassException expected) {
            // expected: either the Object[] component or the depth limit triggers the rejection
        }
    }

    @Test
    public void testExtensionConstructorAllowsAdditionalClass() throws Exception {
        CustomVariable custom = new CustomVariable(99);
        byte[] bytes = serialize(custom);

        // Default strict filter rejects the custom subclass.
        ObjectInputStream strict = new ObjectInputStream(new ByteArrayInputStream(bytes));
        new DefaultMOInput(strict);
        try {
            strict.readObject();
            fail("Expected InvalidClassException for unregistered CustomVariable");
        }
        catch (InvalidClassException expected) {
            // expected
        }

        // When registered via the extension constructor, it is allowed.
        ObjectInputStream extended = new ObjectInputStream(new ByteArrayInputStream(bytes));
        DefaultMOInput in = new DefaultMOInput(extended, Collections.singletonList(CustomVariable.class));
        Object read = in.readVariable();
        assertEquals(custom, read);
    }

    /**
     * A custom serializable {@link Variable} subclass used to verify the extension constructor. Only the
     * (de)serialization behavior is relevant for the test, so the SMI specific methods are minimal.
     */
    public static class CustomVariable extends Integer32 {
        private static final long serialVersionUID = 1L;

        public CustomVariable() {
            super();
        }

        public CustomVariable(int value) {
            super(value);
        }
    }
}
