/*_############################################################################
  _##
  _##  SNMP4J-Agent-DB 3 - MOXodusPersistenceEncryptionTest.java
  _##
  _##  Copyright (C) 2017-2025 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.db;

import jetbrains.exodus.ByteIterable;
import jetbrains.exodus.env.Cursor;
import jetbrains.exodus.env.Environment;
import jetbrains.exodus.env.Environments;
import jetbrains.exodus.env.Store;
import jetbrains.exodus.env.StoreConfig;
import jetbrains.exodus.env.Transaction;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.snmp4j.agent.DefaultMOServer;
import org.snmp4j.agent.DuplicateRegistrationException;
import org.snmp4j.agent.MOServer;
import org.snmp4j.agent.io.ImportMode;
import org.snmp4j.agent.mo.DefaultMOTable;
import org.snmp4j.agent.mo.DefaultMOTableModel;
import org.snmp4j.agent.mo.DefaultMOTableRow;
import org.snmp4j.agent.mo.MOAccessImpl;
import org.snmp4j.agent.mo.MOColumn;
import org.snmp4j.agent.mo.MOScalar;
import org.snmp4j.agent.mo.MOTableIndex;
import org.snmp4j.agent.mo.MOTableSubIndex;
import org.snmp4j.security.SecretOctetString;
import org.snmp4j.smi.Integer32;
import org.snmp4j.smi.OID;
import org.snmp4j.smi.OctetString;
import org.snmp4j.smi.SMIConstants;
import org.snmp4j.smi.Variable;

import java.nio.file.Files;
import java.nio.file.Path;
import java.util.List;

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

/**
 * Tests storing and restoring {@link MOServer} {@link org.snmp4j.agent.ManagedObject}s that hold
 * {@link SecretOctetString} values which need to be encrypted on disk. Both a tabular
 * {@link DefaultMOTable} and a {@link MOScalar} are exercised. The tests verify that:
 * <ul>
 *     <li>store and restore round-trip correctly when the same master password is used,</li>
 *     <li>restore fails (aborts) when a master password was used on store but not on restore,</li>
 *     <li>restore does not yield the original secrets when the master password differs between
 *         store and restore,</li>
 *     <li>the secret values are actually encrypted on disk (the plaintext is not present).</li>
 * </ul>
 *
 * @author Frank Fock
 */
public class MOXodusPersistenceEncryptionTest {

    private static final char[] MASTER_PASSWORD = "correct horse battery staple".toCharArray();
    private static final char[] OTHER_PASSWORD = "a different master password".toCharArray();

    /**
     * Low iteration count keeps PBKDF2 fast for unit testing. The algorithm identifier is stored
     * alongside the ciphertext, so restore uses the embedded parameters regardless of this setting.
     */
    private static final String TEST_ALGORITHM = "PBKDF2WithHmacSHA256:1000:AES/GCM/NoPadding:256";

    private static final OID TABLE_ENTRY_OID = new OID("1.3.6.1.4.1.4976.6.99.1.1");
    private static final int COL_NAME = 2;
    private static final int COL_SECRET = 3;
    private static final OID SCALAR_OID = new OID("1.3.6.1.4.1.4976.6.99.2.0");

    private static final OID ROW1_INDEX = new Integer32(1).toSubIndex(false);
    private static final OID ROW2_INDEX = new Integer32(2).toSubIndex(false);

    private static final OctetString ROW1_NAME = new OctetString("alice");
    private static final OctetString ROW2_NAME = new OctetString("bob");
    private static final byte[] ROW1_SECRET = "alice-secret-key-0123456789".getBytes();
    private static final byte[] ROW2_SECRET = "bob-secret-key-9876543210".getBytes();
    private static final byte[] SCALAR_SECRET = "super-secret-scalar-value".getBytes();

    private static final OctetString SCALAR_PLACEHOLDER = new OctetString("<placeholder>");

    private Path dbDir;

    @Before
    public void setUp() throws Exception {
        dbDir = Files.createTempDirectory("snmp4j-agent-db_encryption-test");
    }

    @After
    public void tearDown() throws Exception {
        // best-effort cleanup of the temporary database directory
        if (dbDir != null) {
            Files.walk(dbDir)
                    .sorted((a, b) -> b.compareTo(a))
                    .forEach(p -> p.toFile().delete());
        }
    }

    /**
     * Holds the managed objects of a freshly built server so the test can inspect them after a load.
     */
    private static final class Fixture {
        final MOServer server;
        final DefaultMOTable<DefaultMOTableRow, MOColumn, DefaultMOTableModel<DefaultMOTableRow>> table;
        final MOScalar<OctetString> scalar;

        Fixture(MOServer server,
                DefaultMOTable<DefaultMOTableRow, MOColumn, DefaultMOTableModel<DefaultMOTableRow>> table,
                MOScalar<OctetString> scalar) {
            this.server = server;
            this.table = table;
            this.scalar = scalar;
        }
    }

    /**
     * Builds a {@link MOServer} with a tabular and a scalar managed object. If {@code withSecretData}
     * is {@code true}, the table is populated with two rows containing {@link SecretOctetString}
     * values and the scalar holds a {@link SecretOctetString}. Otherwise, the table is empty and the
     * scalar holds a non-secret placeholder &mdash; the state into which data is restored.
     */
    @SuppressWarnings({"unchecked", "rawtypes"})
    private Fixture buildFixture(boolean withSecretData) throws DuplicateRegistrationException {
        MOServer server = new DefaultMOServer();

        MOColumn<?>[] columns = new MOColumn<?>[]{
                new MOColumn<>(COL_NAME, SMIConstants.SYNTAX_OCTET_STRING, MOAccessImpl.ACCESS_READ_WRITE),
                new MOColumn<>(COL_SECRET, SMIConstants.SYNTAX_OCTET_STRING, MOAccessImpl.ACCESS_READ_WRITE)
        };
        DefaultMOTable<DefaultMOTableRow, MOColumn, DefaultMOTableModel<DefaultMOTableRow>> table =
                new DefaultMOTable(TABLE_ENTRY_OID,
                        new MOTableIndex(new MOTableSubIndex[]{
                                new MOTableSubIndex(SMIConstants.SYNTAX_INTEGER32)}),
                        columns);

        MOScalar<OctetString> scalar;
        if (withSecretData) {
            table.addRow(new DefaultMOTableRow(ROW1_INDEX,
                    new Variable[]{ROW1_NAME, new SecretOctetString(ROW1_SECRET)}));
            table.addRow(new DefaultMOTableRow(ROW2_INDEX,
                    new Variable[]{ROW2_NAME, new SecretOctetString(ROW2_SECRET)}));
            scalar = new MOScalar<>(SCALAR_OID, MOAccessImpl.ACCESS_READ_WRITE,
                    new SecretOctetString(SCALAR_SECRET));
        } else {
            scalar = new MOScalar<>(SCALAR_OID, MOAccessImpl.ACCESS_READ_WRITE,
                    (OctetString) SCALAR_PLACEHOLDER.clone());
        }

        server.register(table, null);
        server.register(scalar, null);
        return new Fixture(server, table, scalar);
    }

    /**
     * Opens the Xodus environment on the test database directory, builds a server holding the secret
     * data, encrypts and stores it with the given master password, then closes the environment.
     */
    private void storeSecrets(char[] masterPassword) throws DuplicateRegistrationException {
        Environment environment = Environments.newInstance(dbDir.toFile());
        try {
            Fixture fixture = buildFixture(true);
            MOXodusPersistence persistence =
                    new MOXodusPersistence(new MOServer[]{fixture.server}, environment, masterPassword);
            persistence.setDefaultEncryptionAlgorithm(TEST_ALGORITHM);
            persistence.save();
        } finally {
            environment.close();
        }
    }

    /**
     * Opens the Xodus environment on the test database directory, builds an empty server, restores the
     * data with the given master password and returns the fixture for inspection. The environment is
     * closed before returning &mdash; the in-memory managed objects keep their restored state.
     *
     * @param masterPassword
     *    the (non-null) master password to use on restore.
     * @return
     *    the fixture whose managed objects reflect the result of the load operation.
     */
    private Fixture restoreSecrets(char[] masterPassword) throws DuplicateRegistrationException {
        Environment environment = Environments.newInstance(dbDir.toFile());
        try {
            Fixture fixture = buildFixture(false);
            MOXodusPersistence persistence =
                    new MOXodusPersistence(new MOServer[]{fixture.server}, environment, masterPassword);
            persistence.load(ImportMode.restoreChanges);
            return fixture;
        } finally {
            environment.close();
        }
    }

    private static byte[] secretBytes(Fixture fixture, OID rowIndex) {
        DefaultMOTableRow row = fixture.table.getModel().getRow(rowIndex);
        assertNotNull("Expected row " + rowIndex + " to be restored", row);
        Variable value = row.getValue(fixture.table.getColumnIndex(COL_SECRET));
        assertNotNull(value);
        return ((OctetString) value).getValue();
    }

    /**
     * Store and restore using the same master password: both the tabular and scalar secrets must be
     * restored byte-for-byte and the restored values must again be {@link SecretOctetString}s.
     */
    @Test
    public void storeAndRestoreWithSameMasterPassword() throws DuplicateRegistrationException {
        storeSecrets(MASTER_PASSWORD);
        Fixture restored = restoreSecrets(MASTER_PASSWORD);

        // tabular check
        assertEquals(2, restored.table.getModel().getRowCount());
        assertArrayEquals(ROW1_SECRET, secretBytes(restored, ROW1_INDEX));
        assertArrayEquals(ROW2_SECRET, secretBytes(restored, ROW2_INDEX));
        assertTrue("Restored tabular secret should be a SecretOctetString",
                restored.table.getModel().getRow(ROW1_INDEX)
                        .getValue(restored.table.getColumnIndex(COL_SECRET)) instanceof SecretOctetString);
        // the non-secret index column must round-trip too
        assertEquals(ROW1_NAME, restored.table.getModel().getRow(ROW1_INDEX)
                .getValue(restored.table.getColumnIndex(COL_NAME)));

        // scalar check
        assertArrayEquals(SCALAR_SECRET, restored.scalar.getValue().getValue());
        assertTrue("Restored scalar secret should be a SecretOctetString",
                restored.scalar.getValue() instanceof SecretOctetString);
    }

    /**
     * Store with a master password, but restore without one. The loader must abort with a
     * {@link RuntimeException} rather than silently dropping or exposing encrypted data.
     */
    @Test
    public void restoreWithoutPasswordFails() throws DuplicateRegistrationException {
        storeSecrets(MASTER_PASSWORD);

        // This case manages the environment itself rather than going through restoreSecrets(): when
        // load() aborts on the first encrypted value it leaves its read-only transaction open, so a
        // subsequent environment.close() in a finally block would throw "Finish all transactions ..."
        // and mask the RuntimeException we actually want to assert on.
        Environment environment = Environments.newInstance(dbDir.toFile());
        try {
            Fixture fixture = buildFixture(false);
            MOXodusPersistence persistence =
                    new MOXodusPersistence(new MOServer[]{fixture.server}, environment);
            try {
                persistence.load(ImportMode.restoreChanges);
                fail("Expected restore without master password to fail with a RuntimeException");
            } catch (RuntimeException expected) {
                assertNotNull("Expected a descriptive failure message", expected.getMessage());
                assertTrue("Exception should mention the missing master password: " + expected.getMessage(),
                        expected.getMessage().toLowerCase().contains("master password"));
            }
        } finally {
            try {
                environment.close();
            } catch (RuntimeException closeFailure) {
                // A failed load leaves its read-only transaction open; the resulting close() error is a
                // teardown detail unrelated to what this test verifies and is therefore ignored.
            }
        }
    }

    /**
     * Store with one master password and restore with a different one. Decryption of the secret
     * values fails, so the affected instances are not restored: the table stays empty and the scalar
     * keeps its placeholder value. In no case must the original secrets become readable.
     */
    @Test
    public void restoreWithDifferentPasswordDoesNotRecoverSecrets() throws DuplicateRegistrationException {
        storeSecrets(MASTER_PASSWORD);
        Fixture restored = restoreSecrets(OTHER_PASSWORD);

        // The encrypted rows cannot be decrypted -> not imported.
        assertEquals("No tabular rows should be restored with a wrong master password",
                0, restored.table.getModel().getRowCount());

        // The scalar must keep its placeholder and never expose the stored secret.
        assertEquals(SCALAR_PLACEHOLDER, restored.scalar.getValue());
        assertFalse("Scalar must not contain the original secret",
                java.util.Arrays.equals(SCALAR_SECRET, restored.scalar.getValue().getValue()));
    }

    /**
     * Sanity check that the secret values are genuinely encrypted on disk: none of the plaintext
     * secret byte sequences may appear anywhere in the persisted Xodus stores.
     */
    @Test
    public void secretsAreEncryptedOnDisk() throws DuplicateRegistrationException {
        storeSecrets(MASTER_PASSWORD);

        Environment environment = Environments.newInstance(dbDir.toFile());
        try {
            final Transaction txn = environment.beginReadonlyTransaction();
            try {
                List<String> storeNames = environment.getAllStoreNames(txn);
                assertFalse("Expected at least one persisted store", storeNames.isEmpty());
                boolean foundScalar = false;
                boolean foundRow1 = false;
                boolean foundRow2 = false;
                for (String storeName : storeNames) {
                    Store store = environment.openStore(storeName, StoreConfig.WITHOUT_DUPLICATES, txn);
                    try (Cursor cursor = store.openCursor(txn)) {
                        while (cursor.getNext()) {
                            byte[] value = toBytes(cursor.getValue());
                            foundScalar |= contains(value, SCALAR_SECRET);
                            foundRow1 |= contains(value, ROW1_SECRET);
                            foundRow2 |= contains(value, ROW2_SECRET);
                        }
                    }
                }
                assertFalse("Scalar secret found in plaintext on disk", foundScalar);
                assertFalse("Row 1 secret found in plaintext on disk", foundRow1);
                assertFalse("Row 2 secret found in plaintext on disk", foundRow2);
            } finally {
                txn.abort();
            }
        } finally {
            environment.close();
        }
    }

    private static byte[] toBytes(ByteIterable iterable) {
        byte[] unsafe = iterable.getBytesUnsafe();
        byte[] copy = new byte[iterable.getLength()];
        System.arraycopy(unsafe, 0, copy, 0, copy.length);
        return copy;
    }

    /**
     * Returns {@code true} if {@code needle} occurs as a contiguous sub-sequence of {@code haystack}.
     */
    private static boolean contains(byte[] haystack, byte[] needle) {
        if (needle.length == 0 || haystack.length < needle.length) {
            return false;
        }
        outer:
        for (int i = 0; i <= haystack.length - needle.length; i++) {
            for (int j = 0; j < needle.length; j++) {
                if (haystack[i + j] != needle[j]) {
                    continue outer;
                }
            }
            return true;
        }
        return false;
    }
}
