/*_############################################################################
  _## 
  _##  SNMP4J-Agent-DB 3 - SampleAgent.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.sample;

import jetbrains.exodus.env.Environments;
import org.snmp4j.*;
import org.snmp4j.agent.*;
import org.snmp4j.agent.db.MOXodusPersistence;
import org.snmp4j.agent.db.MOXodusPersistenceProvider;
import org.snmp4j.agent.io.ImportMode;
import org.snmp4j.agent.io.MOInputFactory;
import org.snmp4j.agent.io.prop.PropertyMOInput;
import org.snmp4j.agent.mo.*;
import org.snmp4j.agent.mo.snmp.dh.DHKickstartParameters;
import org.snmp4j.agent.mo.snmp.dh.DHKickstartParametersImpl;
import org.snmp4j.log.ConsoleLogFactory;
import org.snmp4j.log.LogAdapter;
import org.snmp4j.log.LogFactory;
import org.snmp4j.log.LogLevel;
import org.snmp4j.mp.MPv3;
import org.snmp4j.security.SecurityProtocols;
import org.snmp4j.smi.*;
import org.snmp4j.transport.DTLSTM;
import org.snmp4j.transport.TransportMappings;
import org.snmp4j.util.ArgumentParser;
import org.snmp4j.util.SnmpConfigurator;
import org.snmp4j.util.ThreadPool;

import javax.crypto.Cipher;
import java.io.*;
import java.security.NoSuchAlgorithmException;
import java.text.ParseException;
import java.util.*;

/**
 * The SampleAgent uses an {@link AgentConfigManager} instance to create a
 * minimal SNMP agent using the configuration defined by
 * {@code SampleAgentConfig.properties} in this package. That properties
 * file defines the initial content of the registered MIB objects of this agent
 * which may differ from the hard coded defaults.
 * <p>
 *     This {@link SampleAgent} uses the {@link MOXodusPersistenceProvider} to persistently
 *     store {@link ManagedObject}s. It assumes that all managed objects that are non-volatile
 *     are implementing the {@link RandomAccessManagedObject} interfaces.
 * <p>
 * The agent uses the console logging to log messages.
 *
 * @author Frank Fock
 * @version 3.7.2
 * @since 3.0
 */
public class SampleAgent implements AgentStateListener<AgentConfigManager> {

    public static final String COMMAND_LINE_OPTIONS = "-c[s{=SampleAgent.cfg}] +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] ";
    public static final String COMMAND_LINE_PARAMS = "#address[s<(udp|tcp|tls|dtls):.*[/[0-9]+]?>] ..";

    static {
        LogFactory.setLogFactory(new ConsoleLogFactory());
        // This enables logging of user credentials (keys and passwords) in plain text to the DEBUG log
        SNMP4JSettings.setSecretLoggingEnabled(true);
    }

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

    protected AgentConfigManager agent;
    protected DefaultMOServer server;
    private final String configFile;
    private final MOXodusPersistence moXodusPersistence;

    protected Properties tableSizeLimits;

    @SuppressWarnings("unchecked")
    public SampleAgent(Map<String, List<Object>> args) {
        LogFactory.getLogFactory().getRootLogger().setLogLevel(LogLevel.ALL);
        //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().setDeadlockPreventionEnabled(true);
        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);
        moXodusPersistence = new MOXodusPersistence(moServers, Environments.newInstance(configFile));

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

        MOInputFactory configurationFactory = null;
        // load initial configuration from properties file only if there is no persistent data for the default context:
        if (!moXodusPersistence.isContextLoadable(null)) {
            configurationFactory = createMOInputFactory(configFilename, ImportMode.restoreChanges);
        }
        this.tableSizeLimits = getTableSizeLimitsProperties(args);
        String dhKickstartInfoPath = (String) ArgumentParser.getFirstValue(args.get("dhks"));
        setupAgent(args, moServers, configurationFactory, args.get("address"), dhKickstartInfoPath);
    }

    /**
     * Read the content of the {@code SampleAgentTableSizeLimits.properties} file.
     * @param args
     *    the command line arguments. If the arguements contain the key "ts" the default file name is replaced by the
     *    key's value.
     * @return
     *    a Properties instance with the table size properties read from the given file.
     */
    protected static Properties getTableSizeLimitsProperties(Map<String, List<Object>> args) {
        InputStream tableSizeLimitsInputStream =
                SampleAgent.class.getResourceAsStream("SampleAgentTableSizeLimits.properties");
        if (args.containsKey("ts")) {
            try {
                tableSizeLimitsInputStream =
                        new FileInputStream((String) ArgumentParser.getValue(args, "ts", 0));
            } catch (FileNotFoundException ex1) {
                ex1.printStackTrace();
            }
        }
        Properties tableSizeLimits = new Properties();
        try {
            tableSizeLimits.load(tableSizeLimitsInputStream);
        } catch (IOException ex) {
            ex.printStackTrace();
        }
        return tableSizeLimits;
    }

    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 SampleAgent(Map<String, List<Object>> args,
                          MOServer[] moServers, MOXodusPersistence xodusPersistence,
                          Properties tableSizeLimits, ImportMode importMode) {
        this.server = new DefaultMOServer();
        this.moXodusPersistence = xodusPersistence;
        List<?> tlsVersions = args.get("tls-version");
        if (tlsVersions != null && (!tlsVersions.isEmpty())) {
            System.setProperty(SnmpConfigurator.P_TLS_VERSION, (String) tlsVersions.get(0));
        }
        String configFilename = null;
        if (args.containsKey("cfg")) {
           configFilename = (String) ArgumentParser.getValue(args, "cfg", 0);
        }
        configFile = (String) (args.get("c")).get(0);
        this.tableSizeLimits = tableSizeLimits;
        String dhKickstartInfoPath = (String) ArgumentParser.getFirstValue(args.get("dhks"));
        setupAgent(args, moServers, createMOInputFactory(configFilename, importMode), args.get("address"), dhKickstartInfoPath);
    }

    protected void setupAgent(Map<String, List<Object>> args,
                              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);
            }
        }
        MOXodusPersistenceProvider moXodusPersistenceProvider = new MOXodusPersistenceProvider(moXodusPersistence);
        OctetString defaultEngineID = new OctetString(MPv3.createLocalEngineID());
        OctetString engineID = moXodusPersistenceProvider.getEngineId(defaultEngineID);
        SnmpConfigurator snmpConfigurator = new SnmpConfigurator(true);
        agent = new AgentConfigManager(
                engineID,
                messageDispatcher,
                null,
                moServers,
                ThreadPool.create("SampleAgent", 3),
                configurationFactory,
                moXodusPersistenceProvider, moXodusPersistenceProvider, null, dhKickstartParameters) {

            @Override
            protected Session createSnmpSession(MessageDispatcher dispatcher) {
                Session session = super.createSnmpSession(dispatcher);
                snmpConfigurator.configure(session, getUsm(), messageDispatcher, args);
                return session;
            }
        };
        agent.addAgentStateListener(this);
    }

    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) {
                // IMPORTANT: Set the server mode for DTLS because for DTLS this cannot be autodetected!
                if (tm instanceof DTLSTM) {
                    ((DTLSTM) tm).setServerEnabled(true);
                }
                md.addTransportMapping(tm);
            } else {
                logger.warn("No transport mapping available for address '" +
                        address + "'.");
            }
        }
    }

    public void run() {
        // initialize agent before registering our own modules
        agent.initialize();
        // switch logging of notifications to log sent notifications instead
        // of logging the original internal notification event:
        //agent.getNotificationLogMIB().setLoggerMode(
        //  NotificationLogMib.Snmp4jNotificationLogModeEnum.sent);
        // this requires sysUpTime to be available.
        // add proxy forwarder
        agent.setupProxyForwarder();

        registerMIBs();
        // apply table size limits
        agent.setTableSizeLimits(tableSizeLimits);
        // register shutdown hook to be able to automatically commit configuration to persistent storage
        agent.registerShutdownHook();
        // now continue agent setup and launch it.
        agent.run();
    }

    /**
     * Get the {@link MOFactory} that creates the various MOs (MIB Objects).
     *
     * @return a {@link DefaultMOFactory} instance by default.
     * @since 1.3.2
     */
    protected MOFactory getFactory() {
        return DefaultMOFactory.getInstance();
    }

    /**
     * Register your own MIB modules in the specified context of the agent.
     * The {@link MOFactory} provided to the {@code Modules} constructor
     * is returned by {@link #getFactory()}.
     */
    protected void registerMIBs() {
    }



    /**
     * Runs a sample 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>
     *
     * @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}.
     */
    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);
            SampleAgent sampleAgent = new SampleAgent(commandLineParameters);
            // Add all available security protocols (e.g. SHA,MD5,DES,AES,3DES,..)
            SecurityProtocols.getInstance().addDefaultProtocols();
            // configure system group:
            // Set system description:
            // sampleAgent.agent.getSysDescr().setValue("My system description".getBytes());
            // Set system OID (= OID of the AGENT-CAPABILITIES statement describing
            // the implemented MIB objects of this agent:
            // sampleAgent.agent.getSysOID().setValue("1.3.1.6.1.4.1....");
            // Set the system services
            // sampleAgent.agent.getSysServices().setValue(72);
            sampleAgent.run();
        } catch (ParseException ex) {
            System.err.println(ex.getMessage());
        }
    }


    /**
     * The agent state has changed to the new state as provided.
     *
     * @param agentConfigManager
     *         the agent's configuration manager. Use this object to access all agent resources, if needed to process this
     *         event.
     * @param newState
     *         the new state of the agent. Although the listener may advance to agent state further, it is not recommended
     *         to do so, because the {@link AgentConfigManager} will do it anyway.
     */
    @Override
    public void agentStateChanged(AgentConfigManager agentConfigManager, AgentState newState) {
        switch (newState.getState()) {
            case AgentState.STATE_INITIALIZED:
                // Optional: Save initialized objects too -> not done by default, because this would be redundant, but
                // in cases, where initial values are determined with high computation effort, using this call might
                // be useful, if there is some logic to avoid second initialization too (not provided here).
                // moXodusPersistence.saveFullDump(null, false);
                moXodusPersistence.registerChangeListenersWithServer(server);
                break;
            case AgentState.STATE_SHUTDOWN:
                moXodusPersistence.unregisterChangeListenersWithServer(server);
                break;
        }
    }
}
