/*_############################################################################
  _## 
  _##  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.audit.*;
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.SecurityLevel;
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.9.2
 * @since 3.0
 */
public class SampleAgent implements AgentStateListener<AgentConfigManager> {

    /** The {@link ArgumentParser} format string describing the supported command line options. */
    public static final String COMMAND_LINE_OPTIONS = "-c[s{=SampleAgent.cfg}] +dhks[s] +u[s] " +
            "+tls-trust-ca[s] +tls-peer-host-id[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] +write-init-config[s] +viewer[s] +admin[s] +x";
    /** The {@link ArgumentParser} format string describing the required command line parameters. */
    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);
    }
    static final AgentAuditPolicy DEFAULT_SAMPLE_AGENT_DB_POLICY = new DefaultAgentAuditPolicy(
            new AgentPolicyProfile(
                    AgentPolicyProfile.ManagedObjectAccess.UNRESTRICTED,
                    AgentPolicyProfile.NetworkAccess.SITE_LOCAL_NETWORK,
                    AgentPolicyProfile.FileAccess.LIMITED_DIRECTORY_ACCESS,
                    SecurityLevel.authNoPriv,
                    new String[]{"SampleAgent.bc","SampleAgent.cfg","SampleAgent-Backup.cfg"},
                    null, null, null
            ), AgentAuditLevel.RESTRICT);

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

    /** The agent configuration manager that controls and runs this sample agent. */
    protected AgentConfigManager agent;
    /** The managed object server that holds the managed objects served by this agent. */
    protected DefaultMOServer server;
    private final String configFile;
    private final MOXodusPersistence moXodusPersistence;

    // Optionally provided named secrets (pass keys) that are injected into the PropertyMOInput as
    // predefined random values "viewer" and "admin" referenced by the configuration properties file.
    private final String viewerPasskey;
    private final String adminPasskey;

    // Optional file name (+write-init-config): if set, a human-readable Markdown overview of the
    // initial MIB configuration is written to this file when the configuration properties are loaded.
    private final String initConfigOverviewFile;

    /** The table size limits applied to the agent's tables, keyed by table OID. */
    protected Properties tableSizeLimits;

    /**
     * The audit policy to be used for the agent.
     */
    private final AgentAuditPolicy auditPolicy;
 
    /**
     * Creates a new {@code SampleAgent} from the parsed command line arguments. Initial configuration is
     * loaded from the configuration properties file only if there is no persistent data for the default context.
     *
     * @param args
     *    the parsed command line arguments as returned by the {@link ArgumentParser}.
     * @param auditPolicy
     *    the audit policy to use for the agent.
     */
    @SuppressWarnings("unchecked")
    public SampleAgent(Map<String, List<Object>> args, AgentAuditPolicy auditPolicy) {
        this.auditPolicy = auditPolicy;
        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);
        }

         /*
           Activate deadlock prevention. This requires few more CPU time and memory consumption during SET request
           and {@link DefaultMOServer#lookup(MOQuery, LockRequest, MOServerLookupEvent, Class)} requests for update
           intention but ensures that SNMP requests and internal updates do not enter a deadlock state (which will
           time out of course but result in a failed request with genErr) when the OIDs locked by a transaction
           are not sorted by their lexicographic order.
         */
        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));

        viewerPasskey = (String) ArgumentParser.getFirstValue(args.get("viewer"));
        adminPasskey = (String) ArgumentParser.getFirstValue(args.get("admin"));

        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 = SampleAgent.getTableSizeLimitsProperties(args, auditPolicy);
        String dhKickstartInfoPath = (String) ArgumentParser.getFirstValue(args.get("dhks"));
        setupAgent(args, moServers, configurationFactory, args.get("address"), dhKickstartInfoPath, auditPolicy);
        initConfigOverviewFile = (String) ArgumentParser.getFirstValue(args.get("write-init-config"));
    }

    /**
     * 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.
     * @param auditPolicy
     *    the audit policy to use or {@code null} if no audit policy is to be used.
     * @return
     *    a Properties instance with the table size properties read from the given file or an empty Properties instance
     *    if the file could not be read.
     */
    protected static Properties getTableSizeLimitsProperties(Map<String, List<Object>> args, AgentAuditPolicy auditPolicy) {
        String configFilename = (String) ArgumentParser.getValue(args, "ts", 0);
        if (auditPolicy == null || configFilename == null || auditPolicy.fileAudit(
                new FileAuditEvent(SampleAgent.class, FileAuditEvent.FileOperation.READ, configFilename))) {
            InputStream tableSizeLimitsInputStream =
                    SampleAgent.class.getResourceAsStream("SampleAgentTableSizeLimits.properties");
            if (configFilename != null) {
                try {
                    tableSizeLimitsInputStream = new FileInputStream(configFilename);
                } catch (FileNotFoundException ex1) {
                    ex1.printStackTrace();
                }
            }
            Properties tableSizeLimits = new Properties();
            try {
                tableSizeLimits.load(tableSizeLimitsInputStream);
            } catch (IOException ex) {
                ex.printStackTrace();
            }
            return tableSizeLimits;
        }
        return new Properties();
    }

    /**
     * Creates a {@link MOInputFactory} that supplies the initial agent configuration from a properties file.
     *
     * @param configFilename
     *    the path to the configuration properties file or {@code null} to use the bundled
     *    {@code SampleAgentConfig.properties} resource.
     * @param importMode
     *    the {@link ImportMode} that controls how the configuration is applied to existing data.
     * @return
     *    a {@link MOInputFactory} providing the parsed configuration.
     */
    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 {
            if (agent.getAuditPolicy().fileAudit(new FileAuditEvent(this,
                    FileAuditEvent.FileOperation.READ, configFile))) {
                configInputStream = new FileInputStream(configFilename);
            }
            else {
                String error = "Config file '" + configFilename + "' not audited: " +
                        "audit policy does not allow reading of file.";
                logger.error(error);
                throw new RuntimeException(error);
            }
            } 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();
        }
        if (initConfigOverviewFile != null) {
            HashMap<String, OctetString> randomValues = new HashMap<>();
            if (viewerPasskey != null) {
                randomValues.put("viewer", new OctetString(viewerPasskey));
            }
            if (adminPasskey != null) {
                randomValues.put("admin", new OctetString(adminPasskey));
            }
            agent.writeInitialConfigOverview(props, initConfigOverviewFile, randomValues);
        }
        configurationFactory = () -> {
            PropertyMOInput propertyMOInput = new PropertyMOInput(props, this.agent, importMode);
            // Inject the optionally provided named secrets (pass keys) so that they are used instead of
            // randomly generated values wherever the configuration references {r}viewer/{r}admin (or
            // {R}viewer/{R}admin). This has to happen before the configuration is read.
            if (viewerPasskey != null) {
                propertyMOInput.setRandomValue("viewer", new OctetString(viewerPasskey));
            }
            if (adminPasskey != null) {
                propertyMOInput.setRandomValue("admin", new OctetString(adminPasskey));
            }
            return propertyMOInput;
        };
        return configurationFactory;
    }

    /**
     * Creates a new {@code SampleAgent} with explicitly provided servers, persistence, and table size limits.
     * This constructor is intended for subclasses and tests that need full control over the agent's collaborators.
     *
     * @param args
     *    the parsed command line arguments as returned by the {@link ArgumentParser}.
     * @param moServers
     *    the managed object servers the agent registers its managed objects with.
     * @param xodusPersistence
     *    the {@link MOXodusPersistence} used to persistently store managed objects.
     * @param tableSizeLimits
     *    the table size limits to apply to the agent's tables.
     * @param importMode
     *    the {@link ImportMode} that controls how the initial configuration is applied.
     * @param auditPolicy
     *    the {@link AgentAuditPolicy} to use for the agent.
     */
    protected SampleAgent(Map<String, List<Object>> args,
                          MOServer[] moServers, MOXodusPersistence xodusPersistence,
                          Properties tableSizeLimits, ImportMode importMode, AgentAuditPolicy auditPolicy) {
        this.auditPolicy = auditPolicy;
        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;
        viewerPasskey = (String) ArgumentParser.getFirstValue(args.get("viewer"));
        adminPasskey = (String) ArgumentParser.getFirstValue(args.get("admin"));
        initConfigOverviewFile = (String) ArgumentParser.getFirstValue(args.get("write-init-config"));
        String dhKickstartInfoPath = (String) ArgumentParser.getFirstValue(args.get("dhks"));
        setupAgent(args, moServers, createMOInputFactory(configFilename, importMode),
                args.get("address"), dhKickstartInfoPath, auditPolicy);
    }

    /**
     * Sets up the {@link AgentConfigManager} including its message dispatcher, listen addresses,
     * optional Diffie-Hellman kickstart parameters, and persistence provider.
     *
     * @param args                 the parsed command line arguments as returned by the {@link ArgumentParser}.
     * @param moServers            the managed object servers the agent registers its managed objects with.
     * @param configurationFactory the {@link MOInputFactory} providing the initial configuration, or {@code null} if none.
     * @param listenAddress        the list of transport addresses the agent should listen on.
     * @param dhKickstartInfoPath  the path to the Diffie-Hellman kickstart parameters file, or {@code null} if none.
     * @param auditPolicy          the {@link AgentAuditPolicy} to use for the agent.
     */
    protected void setupAgent(Map<String, List<Object>> args,
                              MOServer[] moServers, MOInputFactory configurationFactory,
                              List<Object> listenAddress,
                              String dhKickstartInfoPath, AgentAuditPolicy auditPolicy) {
        MessageDispatcher messageDispatcher = new MessageDispatcherImpl();
        addListenAddresses(messageDispatcher, listenAddress);
        Collection<DHKickstartParameters> dhKickstartParameters = Collections.emptyList();
        if (dhKickstartInfoPath != null) {
            if (auditPolicy == null || auditPolicy.fileAudit(
                    new FileAuditEvent(this, FileAuditEvent.FileOperation.READ, dhKickstartInfoPath))) {
                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, auditPolicy);
        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,
                auditPolicy) {

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

    /**
     * Parses the given address strings and adds a matching {@link TransportMapping} for each to the
     * provided {@link MessageDispatcher}.
     *
     * @param md
     *    the message dispatcher the transport mappings are added to.
     * @param addresses
     *    the list of address strings (e.g. {@code udp:0.0.0.0/161}) to listen on.
     */
    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 + "'.");
            }
        }
    }

    /**
     * Initializes the agent, registers the sample MIB modules, applies the table size limits, registers a
     * shutdown hook for persisting the configuration, and finally starts the agent.
     */
    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);
            if (commandLineParameters.containsKey("x")) {
                SNMP4JSettings.setExtensibilityEnabled(true);
            }
            SampleAgent sampleAgent = new SampleAgent(commandLineParameters, DEFAULT_SAMPLE_AGENT_DB_POLICY);
            // 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;
        }
    }
}
