/*_############################################################################
  _## 
  _##  SNMP4J-Agent-DB 3 - MOXodusPersistenceProviderTest.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.env.Environments;
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.snmp.SNMPv2MIB;
import org.snmp4j.agent.mo.snmp.UsmMIB;
import org.snmp4j.mp.MPv3;
import org.snmp4j.security.*;
import org.snmp4j.smi.Integer32;
import org.snmp4j.smi.OID;
import org.snmp4j.smi.OctetString;

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

import static org.junit.Assert.*;

public class MOXodusPersistenceProviderTest {

    public static final String AUTH_12345678 = "auth12345678";
    public static final String AUTH_876543210 = "auth876543210";
    public static final String PRIV_9876543210 = "priv9876543210";
    public static final String PRIV_123445678 = "priv123445678";
    public static final String NO_CONTACT = "no contact";
    public static final String AGENTPP = "AGENTPP";
    private MOServer server;
    private MOXodusPersistence moXodusPersistence;
    private byte[] localEngineID;

    @org.junit.Before
    public void setUp() throws Exception {
        localEngineID = MPv3.createLocalEngineID();
        server = new DefaultMOServer();
        MOServer[] moServers = new MOServer[]{server};
        Path tempDir = Files.createTempDirectory("snmp4j-agent-db_unit-test");
        moXodusPersistence = new MOXodusPersistence(moServers, Environments.newInstance(tempDir.toFile()));
    }

    @org.junit.After
    public void tearDown() throws Exception {

    }

    /**
     * We chane a authPriv and authNoPriv user while leaving a second with the same credentials but different user
     * names the same. After restore, changes should be retained but unchanged entries too.
     * @throws DuplicateRegistrationException
     *    if there is an error in the test.
     */
    @org.junit.Test
    public void restore() throws DuplicateRegistrationException {
        SecurityProtocols.getInstance().addDefaultProtocols();
        OctetString localEngineIDString = new OctetString(localEngineID);
        USM usm = new USM(SecurityProtocols.getInstance(), localEngineIDString, 0);
        UsmMIB usmMIB = new UsmMIB(usm, SecurityProtocols.getInstance());
        usm.addUsmUserListener(usmMIB);
        usmMIB.registerMOs(server, null);

        SNMPv2MIB v2mib = new SNMPv2MIB(new OctetString("test-descr"), new OID("0.0"), new Integer32(2));
        v2mib.registerMOs(server, null);
        doChangesAndRestoreCheck(localEngineIDString, usm, usmMIB, v2mib, true);
    }

    private void doChangesAndRestoreCheck(OctetString localEngineIDString, USM usm, UsmMIB usmMIB, SNMPv2MIB v2mib,
                                          boolean expectObjectsToBeSavedBySaveOperationButNotLater) {
        v2mib.setContact(new OctetString(AGENTPP));

        OctetString userNameAuthPriv = new OctetString("userAuthPriv");
        OctetString userNameAuthNoPriv = new OctetString("userAuthNoPriv");
        OctetString userNameAuthPrivUnchanged = new OctetString("userAuthPrivUnchanged");
        OctetString userNameAuthNoPrivUnchanged = new OctetString("userAuthNoPrivUnchanged");
        OctetString userName = new OctetString("user");

        byte[] keyAuth = new AuthHMAC128SHA224().passwordToKey(new OctetString(AUTH_12345678), localEngineID);
        byte[] keyPriv = new AuthHMAC128SHA224().passwordToKey(new OctetString(PRIV_123445678), localEngineID);
        UsmUser usmUserAuthPriv = new UsmUser(userNameAuthPriv, AuthHMAC128SHA224.ID, new OctetString(keyAuth),
                PrivAES128.ID, new OctetString(keyPriv), localEngineIDString);
        UsmUser usmUserAuthNoPriv = new UsmUser(userNameAuthNoPriv, AuthHMAC128SHA224.ID, new OctetString(keyAuth),
                null, null, localEngineIDString);
        UsmUser usmUserNoAuthNoPriv =
                new UsmUser(userName, null, null, null, null, localEngineIDString);
        UsmUser usmUserAuthPrivUnchanged = new UsmUser(userNameAuthPrivUnchanged, AuthHMAC128SHA224.ID, new OctetString(keyAuth),
                PrivAES128.ID, new OctetString(keyPriv), localEngineIDString);
        UsmUser usmUserAuthNoPrivUnchanged = new UsmUser(userNameAuthNoPrivUnchanged, AuthHMAC128SHA224.ID, new OctetString(keyAuth),
                null, null, localEngineIDString);

        usm.addUser(usmUserAuthNoPrivUnchanged);
        usm.addUser(usmUserAuthNoPriv);
        usm.addUser(usmUserAuthPriv);
        usm.addUser(usmUserNoAuthNoPriv);
        usm.addUser(usmUserAuthPrivUnchanged);

        // Indexes Changed
        OID indexUsmUserAuthNoPriv =
                new OID(new OctetString(localEngineID).toSubIndex(false).getValue(),
                        userNameAuthNoPriv.toSubIndex(false).getValue());
        OID indexUsmUserAuthPriv =
                new OID(new OctetString(localEngineID).toSubIndex(false).getValue(),
                        userNameAuthPriv.toSubIndex(false).getValue());
        // Indexes Unchanged
        OID indexUsmUserAuthNoPrivUnchanged =
                new OID(new OctetString(localEngineID).toSubIndex(false).getValue(),
                        userNameAuthNoPrivUnchanged.toSubIndex(false).getValue());
        OID indexUsmUserAuthPrivUnchanged =
                new OID(new OctetString(localEngineID).toSubIndex(false).getValue(),
                        userNameAuthPrivUnchanged.toSubIndex(false).getValue());

        UsmMIB.UsmTableRow usmTableRow = usmMIB.getUsmUserEntry().getModel().getRow(indexUsmUserAuthNoPriv);
        assertNotNull(usmTableRow);
        assertArrayEquals(keyAuth, ((OctetString)usmTableRow.getValue(UsmMIB.colUsmUserAuthKey)).getValue());

        UsmUserEntry usmUserAuthPrivNew =
                new UsmUserEntry(localEngineID, userNameAuthPriv, AuthHMAC128SHA224.ID,
                        new AuthHMAC128SHA224().passwordToKey(new OctetString(AUTH_876543210), localEngineID),
                PrivAES128.ID, new AuthHMAC128SHA224().passwordToKey(new OctetString(PRIV_9876543210), localEngineID));
        usm.updateUser(usmUserAuthPrivNew);

        UsmUserEntry usmUserAuthNoPrivNew =
                new UsmUserEntry(localEngineID, userNameAuthNoPriv, AuthHMAC128SHA224.ID,
                        new AuthHMAC128SHA224().passwordToKey(new OctetString(AUTH_876543210), localEngineID),
                        null, null);
        usm.updateUser(usmUserAuthNoPrivNew);

        UsmMIB.UsmTableRow usmTableRowNew = usmMIB.getUsmUserEntry().getModel().getRow(indexUsmUserAuthPriv);
        assertNotNull(usmTableRowNew);
        byte[] keyAuthNew = new AuthHMAC128SHA224().passwordToKey(new OctetString(AUTH_876543210), localEngineID);
        byte[] keyPrivNew = new AuthHMAC128SHA224().passwordToKey(new OctetString(PRIV_9876543210), localEngineID);
        assertArrayEquals(keyAuthNew, ((OctetString)usmTableRowNew.getValue(UsmMIB.colUsmUserAuthKey)).getValue());

        moXodusPersistence.save();

        v2mib.setContact(new OctetString(NO_CONTACT));
        assertEquals(new OctetString(NO_CONTACT), v2mib.getContact());
        usmMIB.getUsmUserEntry().removeAll();
        usmTableRowNew = usmMIB.getUsmUserEntry().getModel().getRow(indexUsmUserAuthPriv);
        assertNull(usmTableRowNew);

        moXodusPersistence.load(ImportMode.restoreChanges);

        if (expectObjectsToBeSavedBySaveOperationButNotLater) {
            checkUsmRowsHaveBeenSaved(usmMIB, userNameAuthPriv, userNameAuthNoPriv, userNameAuthPrivUnchanged, userNameAuthNoPrivUnchanged, keyAuth, keyPriv, indexUsmUserAuthNoPriv, indexUsmUserAuthPriv, indexUsmUserAuthNoPrivUnchanged, indexUsmUserAuthPrivUnchanged, keyAuthNew, keyPrivNew);

            // scalar restore check
            assertEquals(new OctetString(AGENTPP), v2mib.getContact());
        }
        else {
            usmTableRowNew = usmMIB.getUsmUserEntry().getModel().getRow(indexUsmUserAuthPriv);
            assertNull(usmTableRowNew);
            // scalar restore check
            assertEquals(new OctetString(NO_CONTACT), v2mib.getContact());
        }
    }

    private void checkUsmRowsHaveBeenSaved(UsmMIB usmMIB, OctetString userNameAuthPriv, OctetString userNameAuthNoPriv, OctetString userNameAuthPrivUnchanged, OctetString userNameAuthNoPrivUnchanged, byte[] keyAuth, byte[] keyPriv, OID indexUsmUserAuthNoPriv, OID indexUsmUserAuthPriv, OID indexUsmUserAuthNoPrivUnchanged, OID indexUsmUserAuthPrivUnchanged, byte[] keyAuthNew, byte[] keyPrivNew) {
        UsmMIB.UsmTableRow usmTableRowNew;
        // Check changed
        usmTableRowNew = usmMIB.getUsmUserEntry().getModel().getRow(indexUsmUserAuthPriv);
        assertNotNull(usmTableRowNew);
        assertArrayEquals(keyAuthNew, ((OctetString)usmTableRowNew.getValue(UsmMIB.colUsmUserAuthKey)).getValue());
        assertArrayEquals(keyPrivNew, ((OctetString)usmTableRowNew.getValue(UsmMIB.colUsmUserPrivKey)).getValue());
        assertEquals(userNameAuthPriv, usmTableRowNew.getValue(UsmMIB.colUsmUserSecurityName));

        usmTableRowNew = usmMIB.getUsmUserEntry().getModel().getRow(indexUsmUserAuthNoPriv);
        assertNotNull(usmTableRowNew);
        assertArrayEquals(keyAuthNew, ((OctetString)usmTableRowNew.getValue(UsmMIB.colUsmUserAuthKey)).getValue());
        assertEquals(userNameAuthNoPriv, ((OctetString)usmTableRowNew.getValue(UsmMIB.colUsmUserSecurityName)));

        // Check unchanged
        usmTableRowNew = usmMIB.getUsmUserEntry().getModel().getRow(indexUsmUserAuthPrivUnchanged);
        assertNotNull(usmTableRowNew);
        assertArrayEquals(keyAuth, ((OctetString)usmTableRowNew.getValue(UsmMIB.colUsmUserAuthKey)).getValue());
        assertArrayEquals(keyPriv, ((OctetString)usmTableRowNew.getValue(UsmMIB.colUsmUserPrivKey)).getValue());
        assertEquals(userNameAuthPrivUnchanged, usmTableRowNew.getValue(UsmMIB.colUsmUserSecurityName));

        usmTableRowNew = usmMIB.getUsmUserEntry().getModel().getRow(indexUsmUserAuthNoPrivUnchanged);
        assertNotNull(usmTableRowNew);
        assertArrayEquals(keyAuth, ((OctetString)usmTableRowNew.getValue(UsmMIB.colUsmUserAuthKey)).getValue());
        assertEquals(userNameAuthNoPrivUnchanged, ((OctetString)usmTableRowNew.getValue(UsmMIB.colUsmUserSecurityName)));
    }

    /**
     * We change a authPriv and authNoPriv user while leaving a second with the same credentials but different user
     * names the same. After restore, changes should be retained but unchanged entries too.
     * @throws DuplicateRegistrationException
     *    if there is an error in the test.
     */
    @org.junit.Test
    public void restoreWithContinuousChangeListening() throws DuplicateRegistrationException {
        SecurityProtocols.getInstance().addDefaultProtocols();
        OctetString localEngineIDString = new OctetString(localEngineID);
        USM usm = new USM(SecurityProtocols.getInstance(), localEngineIDString, 0);
        UsmMIB usmMIB = new UsmMIB(usm, SecurityProtocols.getInstance());
        usm.addUsmUserListener(usmMIB);
        usmMIB.registerMOs(server, null);

        SNMPv2MIB v2mib = new SNMPv2MIB(new OctetString("test-descr"), new OID("0.0"), new Integer32(2));
        v2mib.registerMOs(server, null);

        moXodusPersistence.registerChangeListenersWithServer(server);
        doChangesAndRestoreCheck(localEngineIDString, usm, usmMIB, v2mib, false);
    }

}