/*_############################################################################
  _## 
  _##  SNMP4J-AgentX - TestSubagent.java  
  _## 
  _##  Copyright (C) 2005-2026  Frank Fock (SNMP4J.org)
  _##  
  _##  This program is free software; you can redistribute it and/or modify
  _##  it under the terms of the GNU General Public License version 2 as 
  _##  published by the Free Software Foundation.
  _##
  _##  This program is distributed in the hope that it will be useful,
  _##  but WITHOUT ANY WARRANTY; without even the implied warranty of
  _##  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  _##  GNU General Public License for more details.
  _##
  _##  You should have received a copy of the GNU General Public License
  _##  along with this program; if not, write to the Free Software
  _##  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, 
  _##  MA  02110-1301  USA
  _##  
  _##########################################################################*/

package org.snmp4j.agent.agentx.subagent.test;

import java.io.IOException;
import java.io.InputStream;
import java.util.Collections;
import java.util.List;

import jetbrains.exodus.env.Environments;
import org.snmp4j.agent.*;
import org.snmp4j.agent.agentx.*;
import org.snmp4j.agent.agentx.master.test.TestMasterAgent;
import org.snmp4j.agent.agentx.subagent.*;
import org.snmp4j.agent.audit.AgentAuditLevel;
import org.snmp4j.agent.audit.AgentAuditPolicy;
import org.snmp4j.agent.audit.AgentPolicyProfile;
import org.snmp4j.agent.audit.DefaultAgentAuditPolicy;
import org.snmp4j.agent.db.MOXodusPersistence;
import org.snmp4j.agent.db.MOXodusPersistenceProvider;
import org.snmp4j.agent.io.MOInputFactory;
import org.snmp4j.agent.io.prop.PropertyMOInput;
import org.snmp4j.agent.mo.*;
import org.snmp4j.agent.mo.snmp.SNMPv2MIB.SysUpTimeImpl;
import org.snmp4j.agent.mo.snmp4j.Snmp4jConfigMib;
import org.snmp4j.agent.mo.snmp4j.Snmp4jLogMib;
import org.snmp4j.agent.request.SubRequest;
import org.snmp4j.log.*;
import org.snmp4j.mp.SnmpConstants;
import org.snmp4j.security.SecurityLevel;
import org.snmp4j.smi.*;
import org.snmp4j.transport.ConnectionOrientedTransportMapping;
import org.snmp4j.transport.TransportMappings;
import org.snmp4j.transport.TransportStateEvent;
import org.snmp4j.transport.TransportStateListener;
import org.snmp4j.transport.unix.UnixDomainAddress;
import org.snmp4j.transport.unix.UnixSocketStreamTransportMapping;
import org.snmp4j.util.ArgumentParser;
import org.snmp4j.util.ThreadPool;

import java.util.Map;
import java.util.Properties;

/**
 * The {@code TestSubagent} is an example implementation of a simple
 * AgentX subagent with shared tables and multi context registration.
 *
 * @author Frank Fock
 * @version 4.2.0
 */
public class TestSubagent implements Runnable, TransportStateListener,
        RegistrationCallback {

    static {
        LogFactory.setLogFactory(new ConsoleLogFactory());
        LogFactory.getLogFactory().getRootLogger().setLogLevel(LogLevel.ALL);
    }

    private static final LogAdapter LOGGER =
            LogFactory.getLogger(TestSubagent.class);

    /**
     * The AgentX subagent identifier used when registering this subagent with the master agent.
     */
    public static final OID SUBAGENT_ID = new OID();

    private final AgentXSubagent subagent;
    private final AgentX agentX;
    private final AgentXMessageDispatcher dispatcher;
    private final Address masterAddress;
    private final Address localAddress;
    private AgentXSession<Address> session;
    private int sessionID = 0;

    private final MOServer server;
    private AgentppTestMib agentppTestMib;
    private Snmp4jConfigMib snmp4jConfigMib;
    private Snmp4jLogMib snmp4jLogMib;

    private final SysUpTimeImpl sessionContextUpTime = new SysUpTimeImpl();

    private SubagentXConfigManager configManager;
    private MOXodusPersistence moXodusPersistence;

    static final AgentAuditPolicy SAMPLE_MASTER_AGENT_POLICY = new DefaultAgentAuditPolicy(
            new AgentPolicyProfile(
                    AgentPolicyProfile.ManagedObjectAccess.UNRESTRICTED,
                    AgentPolicyProfile.NetworkAccess.SITE_LOCAL_NETWORK,
                    AgentPolicyProfile.FileAccess.LIMITED_DIRECTORY_ACCESS,
                    SecurityLevel.authNoPriv

            ), AgentAuditLevel.RESTRICT);


    /**
     * Creates a new {@code TestSubagent} that connects to the given AgentX master agent. The constructor
     * sets up the {@link AgentX} instance and its message dispatcher, the {@link MOServer}, the optional
     * persistent storage and configuration factory, and the {@link SubagentXConfigManager} that registers
     * the MIB modules and launches the AgentX sessions.
     *
     * @param masterAddress
     *    the transport address of the AgentX master agent to connect to.
     * @param localAddress
     *    the local transport address used by the subagent.
     * @param configFile
     *    the path to the persistent storage directory, or {@code null} to run without persistence.
     */
    public TestSubagent(Address masterAddress, Address localAddress, String configFile) {
        this.masterAddress = masterAddress;
        this.localAddress = localAddress;
        this.dispatcher = new AgentXMessageDispatcherImpl();
        this.agentX = new AgentX(dispatcher) {
            @Override
            public Map<AgentXSession<?>, Integer> closeAllSessions(byte reason) {
                return subagent.closeAllSessions(reason);
            }
        };
        server = new DefaultMOServer();
        server.addContext(new OctetString());

        MOInputFactory configurationFactory = null;
        InputStream configInputStream = TestSubagent.class.getResourceAsStream("TestSubagentConfig.properties");
        if (configInputStream != null) {
            final Properties props = new Properties();
            try {
                props.load(configInputStream);
            } catch (IOException ex) {
                ex.printStackTrace();
            }
            configurationFactory = () -> new PropertyMOInput(props, configManager);
        }

        MOXodusPersistenceProvider moXodusPersistenceProvider = null;
        if (configFile != null) {
            moXodusPersistence = new MOXodusPersistence(new MOServer[]{server}, Environments.newInstance(configFile));
            moXodusPersistenceProvider = new MOXodusPersistenceProvider(moXodusPersistence, SAMPLE_MASTER_AGENT_POLICY);
        }

        configManager = new SubagentXConfigManager(agentX, dispatcher, new MOServer[]{server},
                ThreadPool.create("AgentXSubAgent", 3),
                configurationFactory, moXodusPersistenceProvider, DefaultMOFactory.getInstance()) {
            /**
             * Register the initialized MIB modules in the specified context of the agent.
             *
             * @param context
             *         the context to register the internal MIB modules. This should be
             *         {@code null} by default.
             *
             * @throws DuplicateRegistrationException
             *         if some of the MIB modules
             *         registration regions conflict with already registered regions.
             */
            @Override
            protected void registerMIBs(OctetString context) throws DuplicateRegistrationException {
                super.registerMIBs(context);
                MOFactory factory = AgentppTestMib.getSharedTableFactory();
                DefaultMOFactory.addSNMPv2TCs(factory);
                agentppTestMib = new AgentppTestMib(factory);
                agentppTestMib.registerMOs(server, null);
            }

            /**
             * Establish AgentX connection to one or more AgentX master agents by implementing this method.
             */
            @Override
            public void launchAgentXSessions() throws IOException {
                unregisterSessionDependent();
                session = new AgentXSession<>(++sessionID);
                int status = subagent.connect(masterAddress, localAddress, session,
                        Collections.singletonList(TestSubagent.this));
                if (status == AgentXProtocol.AGENTX_SUCCESS) {
                    subagent.addAgentCaps(session, new OctetString(),
                            new OID("1.3.6.1.4.1.4976.10.1.1.100.4.1"),
                            new OctetString("AgentX-Test-Subagent"));
                    registerSessionDependent();
                    subagent.registerRegions(session, new OctetString(), null, TestSubagent.this);
                    TimeTicks upTime = new TimeTicks();
                    sessionContextUpTime.setValue(upTime);
                    subagent.setPingDelay(250);
                    subagent.notify(null, SnmpConstants.warmStart,
                            new VariableBinding[]{
                                    new VariableBinding(SnmpConstants.sysDescr,
                                            new OctetString("SNMP4J-AgentX Test-Subagent"))});
                }

            }
        };
        this.subagent = new AgentXSubagent(agentX, SUBAGENT_ID,
                new OctetString("AgentX4J Test agent"), configManager, SAMPLE_MASTER_AGENT_POLICY);
        configManager.registerShutdownHook();
    }

    /**
     * The program entry point that parses the command line arguments, creates the subagent, adds the
     * local transport mapping, and runs the subagent.
     *
     * @param args
     *    the command line arguments.
     */
    public static void main(String[] args) {
        ArgumentParser parser =
                new ArgumentParser("-c[s{=SampleAgent.cfg}] " +
                        "+h +v",
                        "#masteraddress[s{=tcp:0.0.0.0/705}<(unix|tcp):.*[/[0-9]+]?>]");
        Map<String,List<Object>> commandLineParameters = null;
        try {
            commandLineParameters = parser.parse(args);
            if (commandLineParameters.containsKey("h")) {
                printUsage();
                System.exit(0);
            }
            if (commandLineParameters.containsKey("v")) {
                System.out.println("Options: " + commandLineParameters);
            }
            String configFile = (String) (commandLineParameters.get("c")).get(0);
            String agentXAddress =
                    ((String) (commandLineParameters.get("masteraddress")).get(0));
            Address agentXMasterAddress = GenericAddress.parse(agentXAddress);
            if (agentXMasterAddress != null) {
                TransportMappings.getInstance().registerTransportMapping(UnixSocketStreamTransportMapping.class, UnixDomainAddress.class);
                Address localAddress =
                        GenericAddress.newLocalAddress(GenericAddress.getTDomainPrefix(agentXMasterAddress.getClass()));
                TestSubagent sampleAgent = new TestSubagent(agentXMasterAddress, localAddress, configFile);
                sampleAgent.agentX.addTransportMapping(TransportMappings.getInstance().createTransportMapping(localAddress));
                sampleAgent.run();
            }
            else {
                printUsage();
                System.exit(1);
            }
        } catch (ArgumentParser.ArgumentParseException ax) {
            printUsage();
            System.out.println(ax.getMessage());
        } catch (Exception ex) {
            LOGGER.fatal("Caught exception while starting the agent", ex);
            ex.printStackTrace();
        }
    }

    /**
     * Unregisters the session dependent managed objects and removes the session context from the server
     * for the current AgentX session, if any.
     */
    protected void unregisterSessionDependent() {
        if (session != null) {
            OctetString sessionContext = getSessionContext(session.getSessionID());
            server.removeContext(sessionContext);
            if (snmp4jConfigMib != null) {
                snmp4jConfigMib.unregisterMOs(server, sessionContext);
            }
            if (snmp4jLogMib != null) {
                snmp4jLogMib.unregisterMOs(server, sessionContext);
            }
        }
    }

    /**
     * Adds the session context to the server and registers the session dependent managed objects
     * ({@link Snmp4jConfigMib} and {@link Snmp4jLogMib}) for the current AgentX session.
     */
    protected void registerSessionDependent() {
        OctetString sessionContext = getSessionContext(session.getSessionID());
        server.addContext(sessionContext);
        snmp4jConfigMib = new Snmp4jConfigMib(sessionContextUpTime);
        try {
            snmp4jConfigMib.registerMOs(server, sessionContext);
        } catch (DuplicateRegistrationException e) {
            LOGGER.error("Failed to register Snmp4jConfigMib: "+e.getMessage(), e);
        }
        snmp4jLogMib = new Snmp4jLogMib();
        try {
            snmp4jLogMib.registerMOs(server, sessionContext);
        } catch (DuplicateRegistrationException e) {
            LOGGER.error("Failed to register Snmp4jLogMib: "+e.getMessage(), e);
        }
    }

    private static OctetString getSessionContext(int sessionID) {
        return new OctetString("session=" + sessionID);
    }

    public void connectionStateChanged(final TransportStateEvent change) {
        switch (change.getNewState()) {
            case TransportStateEvent.STATE_CONNECTED: {
                LOGGER.info("Connected to "+change.getPeerAddress());
                break;
            }
            case TransportStateEvent.STATE_DISCONNECTED_REMOTELY: {
                // remove session dependent registrations
                // This would be sufficient for AgentX compliant master agents:
                // session.setClosed(true);
                // But we would like to be sure to reset everything in any case:
                subagent.resetConnection((TcpAddress)change.getPeerAddress(), true);
                configManager.getState().setState(AgentState.STATE_SUSPENDED);
                // try to reconnect
                Thread t = new Thread(() -> {
                    Address addr = change.getPeerAddress();
                    boolean reconnected = false;
                    // try to reconnect ten times if we have been disconnected remotely
                    for (int i = 0; i < 10; i++) {
                        try {
                            Thread.sleep(5000);
                            if (subagent.connect(addr, localAddress, session) == AgentXProtocol.AGENTX_SUCCESS) {
                                reconnected = true;
                                configManager.getState().advanceState(AgentState.STATE_RUNNING);
                                // if connected register our MIB objects
                                registerSessionDependent();
                                TimeTicks upTime = new TimeTicks();
                                subagent.registerRegions(session, new OctetString(), upTime,
                                        TestSubagent.this);
                                server.addContext(getSessionContext(session.getSessionID()));
                                sessionContextUpTime.setValue(upTime);
                                break;
                            }
                        } catch (IOException ex) {
                            ex.printStackTrace();
                        } catch (InterruptedException ex) {
                            break;
                        }
                    }
                    if (!reconnected) {
                        configManager.shutdown();
                        System.exit(1);
                    }
                });
                t.start();
                break;
            }
        }
    }

    public void registrationEvent(OctetString context, ManagedObject<?> mo, int status) {
        if (status != AgentXProtocol.AGENTX_SUCCESS) {
            // optionally remove objects from the server,
            // which could not be registered with the master agent here, but
            // that would prevent their registration after a reconnect:
            //      server.unregister(mo, context);
        }
    }

    @Override
    @SuppressWarnings("unchecked")
    public <R extends MOTableRow, C extends MOColumn<?>, M extends MOTableModel<R>> boolean tableRegistrationEvent(
            OctetString context, AgentXSharedMOTable<R,C,M> mo, R row, boolean indexAllocation, int status,
            int retryCount) {
        if ((status != AgentXProtocol.AGENTX_SUCCESS) && (indexAllocation) &&
                ((context == null) || (context.length() == 0)) &&
                (retryCount < 2)) {
            // After an index collision, try to allocate row with index that contains the session ID of the form
            // [<sessionID>]<index>
            if (AgentppTestMib.oidAgentppTestSparseEntry.equals(mo.getOID())) {
                OID failedIndex = row.getIndex();
                int n = failedIndex.get(1) - 48;
                if (mo.removeRow(failedIndex) != null) {
                    do {
                        failedIndex.setValue(new OctetString("[" + session.getSessionID() +
                                "]" +
                                n++).toSubIndex(false).
                                getValue());
                    }
                    while (mo.getModel().containsRow(failedIndex));
                    mo.addRow(row);
                    // retry index allocation
                    return true;
                }
            }
        }
        return false;
    }

    public void unregistrationEvent(OctetString context, ManagedObject<?> mo, int status) {
    }

    @Override
    public <R extends MOTableRow, C extends MOColumn<?>, M extends MOTableModel<R>> void tableUnregistrationEvent(
            OctetString context, AgentXSharedMOTable<R,C,M> mo, R row, boolean indexAllocation, int status) {
    }

    @Override
    public void run() {
        configManager.run();
    }
    private static void printUsage() {
        String[] txt = {
                "Usage: TestMasterAgent [-c <config-path>] [-h] [-v]",
                "                       [-X <masterAddress>] <address1> [<address2> ..]",
                "",
                "where ",
                "  <config-path>        is the directory where persistent MIB data is stored/read.",
                "  -h                   prints this usage help information and exit.",
                "  -v                   print command line parameters.",
                "  <masterAddress>      is the TCP AgentX master agent address of the local",
                "                       host following the format 'tcp:<host>/<port>'.",
                "  <address>            a listen address following the format ",
                "                       'udp|tcp:<host>/<port>', for example udp:0.0.0.0/161",
                ""
        };
        for (String line : txt) {
            System.out.println(line);
        }
    }
}
