/*_############################################################################
  _## 
  _##  SNMP4J-Agent 3 - SnapshotAgent.java  
  _## 
  _##  Copyright (C) 2005-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.test;

import org.snmp4j.*;
import org.snmp4j.agent.*;
import org.snmp4j.agent.example.SampleAgent;
import org.snmp4j.agent.io.DefaultMOPersistenceProvider;
import org.snmp4j.agent.io.ImportMode;
import org.snmp4j.agent.io.MOInputFactory;
import org.snmp4j.agent.io.MOPersistenceProvider;
import org.snmp4j.agent.io.prop.PropertyMOInput;
import org.snmp4j.agent.mo.snmp.*;
import org.snmp4j.agent.mo.snmp.dh.DHKickstartParameters;
import org.snmp4j.agent.mo.snmp.dh.DHKickstartParametersImpl;
import org.snmp4j.agent.request.SubRequest;
import org.snmp4j.cfg.EngineBootsCounterFile;
import org.snmp4j.cfg.EngineBootsProvider;
import org.snmp4j.log.JavaLogFactory;
import org.snmp4j.log.LogLevel;
import org.snmp4j.mp.CounterSupport;
import org.snmp4j.security.*;
import org.snmp4j.smi.OctetString;
import org.snmp4j.smi.OID;
import org.snmp4j.agent.security.MutableVACM;

import java.io.*;

import org.snmp4j.mp.MPv3;
import org.snmp4j.util.ArgumentParser;
import org.snmp4j.util.SnmpConfigurator;
import org.snmp4j.util.ThreadPool;
import org.snmp4j.smi.Variable;
import org.snmp4j.smi.Integer32;

import java.security.NoSuchAlgorithmException;
import java.text.ParseException;
import java.util.*;

import org.snmp4j.log.LogFactory;
import org.snmp4j.log.LogAdapter;
import org.snmp4j.smi.VariableBinding;

import org.snmp4j.agent.mo.ext.StaticMOGroup;
import org.snmp4j.smi.Address;
import org.snmp4j.smi.GenericAddress;
import org.snmp4j.transport.TransportMappings;
import org.snmp4j.agent.mo.snmp.SnmpCommunityMIB.*;

import javax.crypto.Cipher;

/**
 * The SnapshotAgent is provided as an example agent for testing. It reads previously collect variable bindings
 * from a binary file that contains a {@link List} of Java serialized {@link VariableBinding}s.
 * @author Frank Fock
 * @version 3.6.8
 */
public class SnapshotAgent {

    public static final String COMMAND_LINE_OPTIONS = "-c[s{=SnapshotAgent.cfg}] -bc[s{=SnapshotAgent.bc}] +dhks[s] +u[s] " +
            "+tls-trust-ca[s] +tls-peer-id[s] +tls-local-id[s] +tls-version[s{=TLSv1}<(TLSv1|TLSv1.1|TLSv1.2|TLSv1.3)>] +dtls-version[s{=TLSv1.2}<(TLSv1.0|TLSv1.2)>]" +
            "+Djavax.net.ssl.keyStore +Djavax.net.ssl.keyStorePassword " +
            "+Djavax.net.ssl.trustStore +Djavax.net.ssl.trustStorePassword " +
            "+ts[s] +cfg[s] +x ";
    public static final String COMMAND_LINE_PARAMS = "#snapshotFile[s] #address[s<(udp|tcp|tls|dtls):.*[/[0-9]+]?>] ..";

    // initialize logging
    static {
        LogFactory.setLogFactory(new JavaLogFactory());
    }

    private static final LogAdapter logger =
            LogFactory.getLogger(SnapshotAgent.class);

    protected String address;
    protected File snapshot;
    protected List<StaticMOGroup> groups = new ArrayList<StaticMOGroup>();
    protected AgentConfigManager agent;
    protected MOServer server;
    private final String configFile;
    private final File bootCounterFile;

    public SnapshotAgent(Map<String, List<Object>> args) throws IOException {
        LogFactory.getLogFactory().getRootLogger().setLogLevel(LogLevel.ALL);
        SNMP4JSettings.setReportSecurityLevelStrategy(SNMP4JSettings.ReportSecurityLevelStrategy.noAuthNoPrivIfNeeded);
        //Security.setProperty("crypto.policy", "unlimited");
        try {
            logger.info("Max supported AES key length is " + Cipher.getMaxAllowedKeyLength("AES"));
        } catch (NoSuchAlgorithmException e) {
            logger.error("AES privacy not supported by this VM: ", e);
        }

        server = new DefaultMOServer();
        MOServer[] moServers = new MOServer[]{server};
        String configFilename = null;
        if (args.containsKey("cfg")) {
            configFilename = (String) ArgumentParser.getValue(args, "cfg", 0);
        }
        configFile = (String) (args.get("c")).get(0);
        bootCounterFile = new File((String) (args.get("bc")).get(0));

        EngineBootsCounterFile engineBootsCounterFile = new EngineBootsCounterFile(bootCounterFile);
        OctetString ownEngineId = engineBootsCounterFile.getEngineId(new OctetString(MPv3.createLocalEngineID()));

        List<?> tlsVersions = args.get("tls-version");
        if (tlsVersions != null && (!tlsVersions.isEmpty())) {
            System.setProperty(SnmpConfigurator.P_TLS_VERSION, (String) tlsVersions.get(0));
        }

        File snapshot;
        List<Object> snapshotPaths = args.get("snapshotFile");
        snapshot = new File((String) snapshotPaths.get(0));
        if (!snapshot.canRead()) {
            logger.fatal("Snapshot file cannot be read: " + snapshot);
            return;
        }

        MOInputFactory configurationFactory = null;
        // load initial configuration from properties file only if there is no persistent data for the default context:
        configurationFactory = createMOInputFactory(configFilename, ImportMode.restoreChanges);

        String dhKickstartInfoPath = (String) ArgumentParser.getFirstValue(args.get("dhks"));
        DefaultMOPersistenceProvider persistenceProvider = new DefaultMOPersistenceProvider(moServers, configFile);

        setupAgent(args, persistenceProvider, engineBootsCounterFile, ownEngineId, moServers,
                createMOInputFactory(configFilename, ImportMode.replaceCreate), args.get("address"), dhKickstartInfoPath);
        this.snapshot = snapshot;
    }

    protected MOInputFactory createMOInputFactory(String configFilename, ImportMode importMode) {
        MOInputFactory configurationFactory;
        InputStream configInputStream =
                SampleAgent.class.getResourceAsStream("SampleAgentConfig.properties");
        final Properties props = new Properties();
        if (configFilename != null) {
            try {
                configInputStream = new FileInputStream(configFilename);
            } catch (FileNotFoundException ex1) {
                logger.error("Config file '" + configFilename + "' not found: " + ex1.getMessage(), ex1);
                throw new RuntimeException(ex1);
            }
        }
        try {
            props.load(configInputStream);
        } catch (IOException ex) {
            ex.printStackTrace();
        }
        configurationFactory = () -> new PropertyMOInput(props, this.agent, importMode);
        return configurationFactory;
    }

    protected void setupAgent(Map<String, List<Object>> args,
                              MOPersistenceProvider persistenceProvider,
                              EngineBootsProvider engineBootsProvider,
                              OctetString engineID,
                              MOServer[] moServers, MOInputFactory configurationFactory,
                              List<Object> listenAddress,
                              String dhKickstartInfoPath) {
        MessageDispatcher messageDispatcher = new MessageDispatcherImpl();
        addListenAddresses(messageDispatcher, listenAddress);
        Collection<DHKickstartParameters> dhKickstartParameters = Collections.emptyList();
        if (dhKickstartInfoPath != null) {
            File dhKickstartInfoFile = new File(dhKickstartInfoPath);
            if (dhKickstartInfoFile.canRead()) {
                try {
                    Properties kickstartProperties = new Properties();
                    FileInputStream fileInputStream = new FileInputStream(dhKickstartInfoFile);
                    kickstartProperties.load(fileInputStream);
                    fileInputStream.close();
                    dhKickstartParameters =
                            DHKickstartParametersImpl.readFromProperties("org.snmp4j.", kickstartProperties);
                } catch (IOException iox) {
                    logger.error("Failed to load Diffie Hellman kickstart parameters from '" +
                            dhKickstartInfoPath + "': " + iox.getMessage(), iox);
                }
            } else {
                logger.warn("Diffie Hellman kickstart parameters file cannot be read: " + dhKickstartInfoFile);
            }
        }
        SnmpConfigurator snmpConfigurator = new SnmpConfigurator(true);
        agent = new AgentConfigManager(
                engineID,
                messageDispatcher,
                null,
                moServers,
                ThreadPool.create("SampleAgent", 3),
                configurationFactory,
                persistenceProvider, engineBootsProvider, null, dhKickstartParameters) {

            @Override
            protected Session createSnmpSession(MessageDispatcher dispatcher) {
                Session session = super.createSnmpSession(dispatcher);
                snmpConfigurator.configure(session, getUsm(), messageDispatcher, args);
                return session;
            }
        };
        agent.setContext(new SecurityModels(),
                new SecurityProtocols(SecurityProtocols.SecurityProtocolSet.maxCompatibility), new CounterSupport());
    }

    public void run() {
        // initialize agent before registering our own modules
        agent.initialize();
        SNMPv2MIB context1SNMPv2MIB = new SNMPv2MIB(new OctetString(), new OID(), new Integer32(0));
        registerManagedObjects();
        // register shutdown hook to be able to automatically commit configuration to persistent storage
        agent.registerShutdownHook();
        // now continue agent setup and launch it.
        agent.run();
    }

    protected void addListenAddresses(MessageDispatcher md, List<Object> addresses) {
        for (Object addressString : addresses) {
            Address address = GenericAddress.parse(addressString.toString());
            if (address == null) {
                logger.fatal("Could not parse address string '" + addressString + "'");
                return;
            }
            TransportMapping<? extends Address> tm =
                    TransportMappings.getInstance().createTransportMapping(address);
            if (tm != null) {
                md.addTransportMapping(tm);
            } else {
                logger.warn("No transport mapping available for address '" +
                        address + "'.");
            }
        }
    }

    /**
     * Adds initial notification targets and filters.
     *
     * @param targetMIB       the SnmpTargetMIB holding the target configuration.
     * @param notificationMIB the SnmpNotificationMIB holding the notification
     *                        (filter) configuration.
     */
    protected void addNotificationTargets(SnmpTargetMIB targetMIB,
                                          SnmpNotificationMIB notificationMIB) {
    }

    /**
     * Adds all the necessary initial users to the USM.
     *
     * @param usm the USM instance used by this agent.
     */
    protected void addUsmUser(USM usm) {
    }

    /**
     * Adds initial VACM configuration.
     *
     * @param vacm the VacmMIB holding the agent's view configuration.
     */
    protected void addViews(VacmMIB vacm) {
        vacm.addGroup(SecurityModel.SECURITY_MODEL_SNMPv1,
                new OctetString("cpublic"),
                new OctetString("v1v2group"),
                StorageType.nonVolatile);
        vacm.addGroup(SecurityModel.SECURITY_MODEL_SNMPv2c,
                new OctetString("cpublic"),
                new OctetString("v1v2group"),
                StorageType.nonVolatile);
        vacm.addGroup(SecurityModel.SECURITY_MODEL_USM,
                new OctetString("SHADES"),
                new OctetString("v3group"),
                StorageType.nonVolatile);
        vacm.addGroup(SecurityModel.SECURITY_MODEL_USM,
                new OctetString("TEST"),
                new OctetString("v3test"),
                StorageType.nonVolatile);
        vacm.addGroup(SecurityModel.SECURITY_MODEL_USM,
                new OctetString("SHA"),
                new OctetString("v3restricted"),
                StorageType.nonVolatile);
        vacm.addGroup(SecurityModel.SECURITY_MODEL_USM,
                new OctetString("v3notify"),
                new OctetString("v3restricted"),
                StorageType.nonVolatile);

        vacm.addAccess(new OctetString("v1v2group"), new OctetString(),
                SecurityModel.SECURITY_MODEL_ANY,
                SecurityLevel.NOAUTH_NOPRIV,
                MutableVACM.VACM_MATCH_EXACT,
                new OctetString("fullReadView"),
                new OctetString("fullWriteView"),
                new OctetString("fullNotifyView"),
                StorageType.nonVolatile);
        vacm.addAccess(new OctetString("v3group"), new OctetString(),
                SecurityModel.SECURITY_MODEL_USM,
                SecurityLevel.AUTH_PRIV,
                MutableVACM.VACM_MATCH_EXACT,
                new OctetString("fullReadView"),
                new OctetString("fullWriteView"),
                new OctetString("fullNotifyView"),
                StorageType.nonVolatile);
        vacm.addAccess(new OctetString("v3restricted"), new OctetString(),
                SecurityModel.SECURITY_MODEL_USM,
                SecurityLevel.NOAUTH_NOPRIV,
                MutableVACM.VACM_MATCH_EXACT,
                new OctetString("restrictedReadView"),
                new OctetString("restrictedWriteView"),
                new OctetString("restrictedNotifyView"),
                StorageType.nonVolatile);
        vacm.addAccess(new OctetString("v3test"), new OctetString(),
                SecurityModel.SECURITY_MODEL_USM,
                SecurityLevel.AUTH_PRIV,
                MutableVACM.VACM_MATCH_EXACT,
                new OctetString("testReadView"),
                new OctetString("testWriteView"),
                new OctetString("testNotifyView"),
                StorageType.nonVolatile);

        vacm.addViewTreeFamily(new OctetString("fullReadView"), new OID("1.3"),
                new OctetString(), VacmMIB.vacmViewIncluded,
                StorageType.nonVolatile);
        vacm.addViewTreeFamily(new OctetString("fullWriteView"), new OID("1.3"),
                new OctetString(), VacmMIB.vacmViewIncluded,
                StorageType.nonVolatile);
        vacm.addViewTreeFamily(new OctetString("fullNotifyView"), new OID("1.3"),
                new OctetString(), VacmMIB.vacmViewIncluded,
                StorageType.nonVolatile);

        vacm.addViewTreeFamily(new OctetString("restrictedReadView"),
                new OID("1.3.6.1.2"),
                new OctetString(), VacmMIB.vacmViewIncluded,
                StorageType.nonVolatile);
        vacm.addViewTreeFamily(new OctetString("restrictedWriteView"),
                new OID("1.3.6.1.2.1"),
                new OctetString(),
                VacmMIB.vacmViewIncluded,
                StorageType.nonVolatile);
        vacm.addViewTreeFamily(new OctetString("restrictedNotifyView"),
                new OID("1.3.6.1.2"),
                new OctetString(), VacmMIB.vacmViewIncluded,
                StorageType.nonVolatile);
        vacm.addViewTreeFamily(new OctetString("restrictedNotifyView"),
                new OID("1.3.6.1.6.3.1"),
                new OctetString(), VacmMIB.vacmViewIncluded,
                StorageType.nonVolatile);

        vacm.addViewTreeFamily(new OctetString("testReadView"),
                new OID("1.3.6.1.2"),
                new OctetString(), VacmMIB.vacmViewIncluded,
                StorageType.nonVolatile);
        vacm.addViewTreeFamily(new OctetString("testReadView"),
                new OID("1.3.6.1.2.1.1"),
                new OctetString(), VacmMIB.vacmViewExcluded,
                StorageType.nonVolatile);
        vacm.addViewTreeFamily(new OctetString("testWriteView"),
                new OID("1.3.6.1.2.1"),
                new OctetString(),
                VacmMIB.vacmViewIncluded,
                StorageType.nonVolatile);
        vacm.addViewTreeFamily(new OctetString("testNotifyView"),
                new OID("1.3.6.1.2"),
                new OctetString(), VacmMIB.vacmViewIncluded,
                StorageType.nonVolatile);

    }

    /**
     * Register additional managed objects at the agent's server.
     */
    protected void registerManagedObjects() {
        FileInputStream fis = null;
        try {
            fis = new FileInputStream(snapshot);
            ObjectInputStream ois = new ObjectInputStream(fis);
            @SuppressWarnings("unchecked")
            List<VariableBinding> l = (List<VariableBinding>) ois.readObject();
            ois.close();
            logger.info("Snapshot file '" + snapshot +
                    "' contains " + l.size() + " objects.");
            OctetString ctx = new OctetString();
            SortedMap<OID, OID> rootsCopy = getSubtreeRoots(l);
            logger.info("Identified the following sub-trees " + rootsCopy);
            for (OID root : rootsCopy.keySet()) {
                ArrayList<VariableBinding> subtree = new ArrayList<VariableBinding>();
                for (VariableBinding vb : l) {
                    if (vb.getOid().size() >= root.size()) {
                        if (vb.getOid().leftMostCompare(root.size(), root) == 0) {
                            subtree.add(vb);
                        }
                    }
                }
                StaticMOGroup group =
                        new StaticMOGroup(root, subtree.toArray(new VariableBinding[0]));
                DefaultMOContextScope scope = new DefaultMOContextScope(ctx,
                        root, true, root.nextPeer(), false);
                MOServerLookupEvent lookupEvent = new MOServerLookupEvent(this, null,
                        new DefaultMOQuery(scope, false), MOServerLookupEvent.IntendedUse.register, true);
                ManagedObject<?> mo = server.lookup(lookupEvent.getQuery(), null, lookupEvent);
                if (mo != null) {
                    logger.warn("Could not register snapshot subtree '" + root +
                            "' with " + subtree + " because ManagedObject " + mo +
                            " is already registered");
                    for (VariableBinding vb : subtree) {
                        group = new StaticMOGroup(vb.getOid(), new VariableBinding[]{vb});
                        scope = new DefaultMOContextScope(ctx, vb.getOid(), true, vb.getOid().nextPeer(), false);
                        lookupEvent = new MOServerLookupEvent(this, null,
                                new DefaultMOQuery(scope, false),
                                MOServerLookupEvent.IntendedUse.register, true);
                        mo = server.lookup(lookupEvent.getQuery(), null, lookupEvent);
                        if (mo != null) {
                            logger.warn("Could not register single OID at " + vb.getOid() +
                                    " because ManagedObject " + mo + " is already registered.");
                        } else {
                            groups.add(group);
                            server.register(group, null);
                            logger.info("Registered snapshot subtree '" + root + "' with " + vb);
                            lookupEvent.completedUse(group);
                        }
                    }
                } else {
                    groups.add(group);
                    server.register(group, null);
                    logger.info("Registered snapshot subtree '" + root +
                            "' with " + subtree);
                    lookupEvent.completedUse(group);
                }
            }
        } catch (Exception ex) {
            logger.error("Error while reading snapshot file '" + snapshot + "':" +
                    ex.getMessage(), ex);
        }
    }

    static SortedMap<OID, OID> getSubtreeRoots(List<VariableBinding> l) {
        SortedMap<OID, OID> roots = new TreeMap<>();
        for (VariableBinding vb : l) {
            if (roots.isEmpty()) {
                OID rootCandidate = vb.getOid().trim();
                if (rootCandidate.size() >= 2) {
                    roots.put(rootCandidate, rootCandidate);
                }
            }
            else {
                boolean rootFound = false;
                for (OID rootCandidate : roots.keySet()) {
                    if (vb.getOid().startsWith(rootCandidate)) {
                        rootFound = true;
                        break;
                    }
                    else {
                        OID newRootCandidate = vb.getOid().trim();
                        if (rootCandidate.startsWith(newRootCandidate)) {
                            roots.remove(rootCandidate);
                            roots.put(newRootCandidate, newRootCandidate);
                            rootFound= true;
                        }
                    }
                }
                if (!rootFound) {
                    OID rootCandidate = vb.getOid().trim();
                    if (rootCandidate.size() >= 2) {
                        roots.put(rootCandidate, rootCandidate);
                    }
                }
            }
        }
        logger.info("Identified the following sub-tree candidates: " + roots);
        SortedMap<OID, OID> rootsCopy = new TreeMap<>();
        for (OID k : roots.keySet()) {
            if (k.size() > 1) {
                OID sk = new OID(k.getValue(), 0, k.size() - 1);
                while ((sk.size() > 0) && (roots.get(sk) == null)) {
                    sk.trim(1);
                }
                if (sk.size() == 0) {
                    rootsCopy.put(k, k);
                }
            }
        }
        return rootsCopy;
    }

    /**
     * Unregister additional managed objects from the agent's server.
     */
    protected void unregisterManagedObjects() {
        for (StaticMOGroup mo : groups) {
            server.unregister(mo, null);
        }
    }

    /**
     * Runs a snapshot agent with a default configuration defined by {@code SampleAgentConfig.properties}. A sample
     * command line is:
     * <pre>
     * -c SampleAgent.cfg -bc SampleAgent.bc udp:127.0.0.1/4700 tcp:127.0.0.1/4700
     * </pre>
     * Separated by || options for more than one agent in the same Java process can be given.
     *
     * @param args the command line arguments defining at least the listen addresses. The format is {@code
     *             -c[s{=SampleAgent.cfg}] -bc[s{=SampleAgent.bc}] +ts[s] +cfg[s] #address[s&lt;(udp|tcp|tls):.*[/[0-9]+]?&gt;]
     *             ..}. For the format description see {@link ArgumentParser}. Multiple agent configurations can be given
     *             separated by ||.
     */
    public static void main(String[] args) {
        ArgumentParser parser =
                new ArgumentParser(COMMAND_LINE_OPTIONS, COMMAND_LINE_PARAMS);
        Map<String, List<Object>> commandLineParameters;
        try {
            commandLineParameters = parser.parse(args);
            if (commandLineParameters.containsKey("x")) {
                SNMP4JSettings.setExtensibilityEnabled(true);
            }
            SnapshotAgent snapshotAgent = new SnapshotAgent(commandLineParameters);
            // Add all available security protocols (e.g. SHA,MD5,DES,AES,3DES,..)
            SecurityProtocols.getInstance().addDefaultProtocols();
            snapshotAgent.run();
        } catch (ParseException | IOException ex) {
            ex.printStackTrace();
        }
    }
}