# -*- coding: utf-8 -*-

# Copyright the NTPsec project contributors
#
# SPDX-License-Identifier: BSD-2-Clause

from __future__ import print_function, division

import select
import time
import sys

try:
    import ntp.util
    import ntp.agentx_packet
    ax = ntp.agentx_packet
except ImportError as e:
    sys.stderr.write(
        "AgentX: can't find Python AgentX Packet library.\n")
    sys.stderr.write("%s\n" % e)
    sys.exit(1)


defaultTimeout = 30
pingTime = 60


def gen_next(generator):
    if str is bytes:  # Python 2
        return generator.next()
    else:  # Python 3
        return next(generator)


class MIBControl:
    def __init__(self, oidTree=None, mibRoot=(), rangeSubid=0, upperBound=None,
                 mibContext=None):
        self.oidTree = {}  # contains callbacks for the MIB
        if oidTree is not None:
            self.oidTree = oidTree
        # The undo system is only for the last operation
        self.inSetP = False  # Are we currently in the set procedure?
        self.setVarbinds = []  # Varbind of the current set operation
        self.setHandlers = []  # Handlers for commit/undo/cleanup of set
        self.setUndoData = []  # Previous values for undoing
        self.mibRoot = mibRoot
        self.rangeSubid = rangeSubid
        self.upperBound = upperBound
        self.context = mibContext

    def mib_rootOID(self):
        return self.mibRoot

    def mib_rangeSubid(self):
        return self.rangeSubid

    def mib_upperBound(self):
        return self.upperBound

    def mib_context(self):
        return self.context

    def addNode(self, oid, reader=None, writer=None, dynamic=None):
        if isinstance(oid, ax.OID):  # get it in a mungable format
            oid = tuple(oid.subids)
        # dynamic is the generator for tables
        currentLevel = self.oidTree
        remainingOID = oid
        while True:
            node, remainingOID = ntp.util.slicedata(remainingOID, 1)
            node = node[0]
            if node not in currentLevel.keys():
                currentLevel[node] = {"reader": None, "writer": None,
                                      "subids": None}
            if not remainingOID:  # We have reached the target node
                currentLevel[node]["reader"] = reader
                currentLevel[node]["writer"] = writer
                if dynamic is not None:
                    # can't be both dynamic and non-dynamic
                    currentLevel[node]["subids"] = dynamic
                return
            else:
                if currentLevel[node]["subids"] is None:
                    currentLevel[node]["subids"] = {}
                currentLevel = currentLevel[node]["subids"]

    def getOID_core(self, nextP, searchoid, returnGenerator=False):
        gen = walkMIBTree(self.oidTree, self.mibRoot)
        while True:
            try:
                oid, reader, writer = gen_next(gen)
                if nextP:  # GetNext
                    # For getnext any OID greater than the start qualifies
                    oidhit = (oid > searchoid)
                else:  # Get
                    # For get we need a *specific* OID
                    oidhit = (oid.subids == searchoid.subids)
                if oidhit and (reader is not None):
                    # We only return OIDs that have a minimal implementation
                    # walkMIBTree handles the generation of dynamic trees
                    if returnGenerator:
                        return oid, reader, writer, gen
                    else:
                        return oid, reader, writer
            except StopIteration:  # Couldn't find anything in the tree
                if returnGenerator:
                    return None, None, None, None
                else:
                    return None, None, None

    # These exist instead of just using getOID_core so semantics are clearer
    def getOID(self, searchoid, returnGenerator=False):
        "Get the requested OID"
        return self.getOID_core(False, searchoid, returnGenerator)

    def getNextOID(self, searchoid, returnGenerator=False):
        "Get the next lexicographical OID"
        return self.getOID_core(True, searchoid, returnGenerator)

    def getOIDsInRange(self, oidrange, firstOnly=False):
        "Get a list of every (optionally the first) OID in a range"
        oids = []
        gen = walkMIBTree(self.oidTree, self.mibRoot)
        # Find the first OID
        while True:
            try:
                oid, reader, writer = gen_next(gen)
                if reader is None:
                    continue  # skip unimplemented OIDs
                elif oid.subids == oidrange.start.subids:
                    # ok, found the start, do we need to skip it?
                    if oidrange.start.include:
                        oids.append((oid, reader, writer))
                        break
                    else:
                        continue
                elif oid > oidrange.start:
                    # If we are here it means we hit the start but skipped
                    if not oidrange.end.isNull() and oid >= oidrange.end:
                        # We fell off the range
                        return []
                    oids.append((oid, reader, writer))
                    break
            except StopIteration:
                # Couldn't find *anything*
                return []
        if firstOnly:
            return oids
        # Start filling in the rest of the range
        while True:
            try:
                oid, reader, writer = gen_next(gen)
                if reader is None:
                    continue  # skip unimplemented OIDs
                elif not oidrange.end.isNull() and oid >= oidrange.end:
                    break  # past the end of a bounded range
                else:
                    oids.append((oid, reader, writer))
            except StopIteration:
                break  # We have run off the end of the MIB
        return oids


class PacketControl:
    def __init__(self, sock, dbase, spinGap=0.001, timeout=defaultTimeout,
                 logfp=None, debug=10000):
        self.log = (lambda txt, dbg: ntp.util.dolog(logfp, txt, debug, dbg))
        # take a pre-made socket instead of making our own so that
        # PacketControl doesn't have to know or care about implementation
        self.socket = sock
        self.spinGap = spinGap  # sleep() time on each loop
        # indexed on: (session_id, transaction_id, packet_id)
        # contains: (timeout, packet class)
        self.packetLog = {}  # Sent packets kept until response is received
        self.loopCallback = None  # called each loop in runforever mode
        self.database = dbase  # class for handling data requests
        self.receivedData = b""  # buffer for data from incomplete packets
        self.receivedPackets = []  # use as FIFO
        self.timeout = timeout
        self.sessionID = None  # need this for all packets
        self.highestTransactionID = 0  # used for exchanges we start
        self.lastReception = None
        self.stillConnected = False
        # indexed on pdu code
        self.pduHandlers = {ax.PDU_GET: self.handle_GetPDU,
                            ax.PDU_GET_NEXT: self.handle_GetNextPDU,
                            ax.PDU_GET_BULK: self.handle_GetBulkPDU,
                            ax.PDU_TEST_SET: self.handle_TestSetPDU,
                            ax.PDU_COMMIT_SET: self.handle_CommitSetPDU,
                            ax.PDU_UNDO_SET: self.handle_UndoSetPDU,
                            ax.PDU_CLEANUP_SET: self.handle_CleanupSetPDU,
                            ax.PDU_RESPONSE: self.handle_ResponsePDU}

    def mainloop(self, runforever):
        if self.stillConnected is not True:
            return False
        if runforever:
            while self.stillConnected:
                self._doloop()
                if self.loopCallback is not None:
                    self.loopCallback(self)
                time.sleep(self.spinGap)
        else:
            self._doloop()
        return self.stillConnected

    def _doloop(self):
        # loop body split out to separate the one-shot/run-forever switches
        # from the actual logic
        self.packetEater()
        while self.receivedPackets:
            packet = self.receivedPackets.pop(0)
            if packet.sessionID != self.sessionID:
                self.log(
                    "Received packet with incorrect session ID: %s" % packet,
                    3)
                resp = ax.ResponsePDU(True, packet.sessionID,
                                      packet.transactionID, packet.packetID,
                                      0, ax.RSPERR_NOT_OPEN, 0)
                self.sendPacket(resp, False)
                continue
            ptype = packet.pduType
            if ptype in self.pduHandlers:
                self.pduHandlers[ptype](packet)
            else:
                self.log("Dropping packet type %i, not implemented" % ptype,
                         2)
        self.checkResponses()
        if self.lastReception is not None:
            currentTime = time.time()
            if (currentTime - self.lastReception) > pingTime:
                self.sendPing()

    def initNewSession(self):
        self.log("Initializing new session...", 3)
        # We already have a connection, need to open a session.
        openpkt = ax.OpenPDU(True, 23, 0, 0, self.timeout, (),
                             "NTPsec SNMP subagent")
        self.sendPacket(openpkt, False)
        response = self.waitForResponse(openpkt, True)
        self.sessionID = response.sessionID
        # Register the tree
        register = ax.RegisterPDU(True, self.sessionID, 1, 1, self.timeout, 1,
                                  self.database.mib_rootOID(),
                                  self.database.mib_rangeSubid(),
                                  self.database.mib_upperBound(),
                                  self.database.mib_context())
        self.sendPacket(register, False)
        self.waitForResponse(register)
        self.stillConnected = True

    def waitForResponse(self, opkt, ignoreSID=False):
        "Wait for a response to a specific packet, dropping everything else"
        while True:
            self.packetEater()
            while self.receivedPackets:
                packet = self.receivedPackets.pop(0)
                if packet.__class__ != ax.ResponsePDU:
                    continue
                haveit = (opkt.transactionID == packet.transactionID) and \
                         (opkt.packetID == packet.packetID)
                if not ignoreSID:
                    haveit = haveit and (opkt.sessionID == packet.sessionID)
                if haveit:
                    self.log("Received waited for response", 4)
                    return packet
            time.sleep(self.spinGap)

    def checkResponses(self):
        "Check for expected responses that have timed out"
        currentTime = time.time()
        for key in list(self.packetLog.keys()):
            expiration, originalPkt, callback = self.packetLog[key]
            if currentTime > expiration:
                if callback is not None:
                    callback(None, originalPkt)
                del self.packetLog[key]

    def packetEater(self):
        "Slurps data from the input buffer and tries to parse packets from it"
        self.pollSocket()
        while True:
            datalen = len(self.receivedData)
            if datalen < 20:
                return None  # We don't even have a packet header, bail
            try:
                pkt, fullPkt, extraData = ax.decode_packet(self.receivedData)
                if not fullPkt:
                    return None
                self.receivedData = extraData
                self.receivedPackets.append(pkt)
                if pkt.transactionID > self.highestTransactionID:
                    self.highestTransactionID = pkt.transactionID
                self.log("Received a full packet: %s" % repr(pkt), 4)
            except (ax.ParseVersionError, ax.ParsePDUTypeError,
                    ax.ParseError) as e:
                if e.header["type"] != ax.PDU_RESPONSE:
                    # Response errors are silently dropped, per RFC
                    # Everything else sends an error response
                    self.sendErrorResponse(e.header, ax.RSPERR_PARSE_ERROR, 0)
                # *Hopefully* the packet length was correct.....
                #  if not, all packets will be scrambled. Maybe dump the
                #  whole buffer if too many failures in a row?
                self.receivedData = e.remainingData

    def sendPacket(self, packet, expectsReply, replyTimeout=defaultTimeout,
                   callback=None):
        encoded = packet.encode()
        self.log("Sending packet (with reply: %s): %s" % (expectsReply,
                                                          repr(packet)), 4)
        self.socket.sendall(encoded)
        if expectsReply:
            index = (packet.sessionID,
                     packet.transactionID,
                     packet.packetID)
            self.packetLog[index] = (replyTimeout, packet, callback)

    def sendPing(self):
        # DUMMY packetID, does this need to change? or does the pktID only
        # count relative to a given transaction ID?
        tid = self.highestTransactionID + 5  # +5 to avoid collisions
        self.highestTransactionID = tid
        pkt = ax.PingPDU(True, self.sessionID, tid, 1)

        def callback(resp, orig):
            if resp is None:  # Timed out. Need to restart the session.
                # Er, problem: Can't handle reconnect from inside PacketControl
                self.stillConnected = False
        self.sendPacket(pkt, True, callback=callback)

    def sendNotify(self, varbinds, context=None):
        # DUMMY packetID, does this need to change? or does the pktID only
        # count relative to a given transaction ID?
        tid = self.highestTransactionID + 5  # +5 to avoid collisions
        self.highestTransactionID = tid
        pkt = ax.NotifyPDU(True, self.sessionID, tid, 1, varbinds, context)

        def resendNotify(pkt, orig):
            if pkt is None:
                self.sendPacket(orig, True, callback=resendNotify)
        self.sendPacket(pkt, True, resendNotify)

    def sendErrorResponse(self, errorHeader, errorType, errorIndex):
        err = ax.ResponsePDU(errorHeader["flags"]["bigEndian"],
                             errorHeader["session_id"],
                             errorHeader["transaction_id"],
                             errorHeader["packet_id"],
                             0, errorType, errorIndex)
        self.sendPacket(err, False)

    def pollSocket(self):
        "Reads all currently available data from the socket, non-blocking"
        data = b""
        while True:
            tmp = select.select([self.socket], [], [], 0)[0]
            if not tmp:  # No socket, means no data available
                break
            tmp = tmp[0]
            newdata = tmp.recv(4096)  # Arbitrary value
            if newdata:
                self.log("Received data: %s" % repr(newdata), 5)
                data += newdata
                self.lastReception = time.time()
            else:
                break
        self.receivedData += data

    # ==========================
    # Packet handlers start here
    # ==========================

    def handle_GetPDU(self, packet):
        binds = []
        for oidr in packet.oidranges:
            target = oidr.start
            oid, reader, _ = self.database.getOID(target)
            if (oid != target) or (reader is None):
                # This OID must not be implemented yet.
                binds.append(ax.Varbind(ax.VALUE_NO_SUCH_OBJECT, target))
            else:
                vbind = reader(oid)
                if vbind is None:  # No data available.
                    # I am not certain that this is the correct response
                    # when no data is available. snmpwalk appears to stop
                    # calling a particular sub-agent when it gets to a NULL.
                    binds.append(ax.Varbind(ax.VALUE_NULL, target))
                else:
                    binds.append(vbind)
            # There should also be a situation that leads to noSuchInstance
            #  but I do not understand the requirements for that
        # TODO: Need to implement genError
        resp = ax.ResponsePDU(True, self.sessionID, packet.transactionID,
                              packet.packetID, 0, ax.ERR_NOERROR, 0, binds)
        self.sendPacket(resp, False)

    def handle_GetNextPDU(self, packet):
        binds = []
        for oidr in packet.oidranges:
            while True:
                oids = self.database.getOIDsInRange(oidr, True)
                if not oids:  # Nothing found
                    binds.append(ax.Varbind(ax.VALUE_END_OF_MIB_VIEW,
                                            oidr.start))
                    break
                else:
                    oid, reader, _ = oids[0]
                    vbind = reader(oid)
                    if vbind is None:  # No data available
                        # Re-do search for this OID range, starting from just
                        # after the current location
                        oidr = ax.SearchRange(oid, oidr.end, False)
                        continue
                    else:
                        binds.append(vbind)
                        break
        # TODO: Need to implement genError
        resp = ax.ResponsePDU(True, self.sessionID, packet.transactionID,
                              packet.packetID, 0, ax.ERR_NOERROR, 0, binds)
        self.sendPacket(resp, False)

    def handle_GetBulkPDU(self, packet):
        binds = []
        nonreps = packet.oidranges[:packet.nonReps]
        repeats = packet.oidranges[packet.nonReps:]
        # Handle non-repeats
        for oidr in nonreps:
            oids = self.database.getOIDsInRange(oidr, True)
            if not oids:  # Nothing found
                binds.append(ax.Varbind(ax.VALUE_END_OF_MIB_VIEW, oidr.start))
            else:
                oid, reader, _ = oids[0]
                binds.append(reader(oid))
        # Handle repeaters
        for oidr in repeats:
            oids = self.database.getOIDsInRange(oidr)
            if not oids:  # Nothing found
                binds.append(ax.Varbind(ax.VALUE_END_OF_MIB_VIEW, oidr.start))
            else:
                for oid, reader, _ in oids[:packet.maxReps]:
                    binds.append(reader(oid))
        resp = ax.ResponsePDU(True, self.sessionID, packet.transactionID,
                              packet.packetID, 0, ax.ERR_NOERROR, 0, binds)
        self.sendPacket(resp, False)

    def handle_TestSetPDU(self, packet):  # WIP / TODO
        # Be advised: MOST OF THE VALIDATION IS DUMMY CODE OR DOESN'T EXIST
        # According to the RFC this is one of the most demanding parts and
        #  *has* to be gotten right
        if self.database.inSetP:
            pass  # Is this an error?
        # if (inSetP) is an error these will go in an else block
        self.database.inSetP = True
        self.database.setVarbinds = []
        self.database.setHandlers = []
        self.database.setUndoData = []
        error = None
        for bindIndex in range(len(packet.varbinds)):
            varbind = packet.varbinds[bindIndex]
            # Find an OID, then validate it
            oid, reader, writer = self.database.getOID(varbind.oid)
            if oid is None:  # doesn't exist, can we create it?
                # DUMMY, assume we can't create anything
                error = ax.ERR_NO_ACCESS
                break
            elif writer is None:  # exists, writing not implemented
                error = ax.ERR_NOT_WRITABLE
                break
            # Ok, we have an existing or new OID, assemble the orders
            # If we created a new bind undoData is None, must delete it
            undoData = reader(oid)
            error = writer("test", varbind)
            if error != ax.ERR_NOERROR:
                break
            self.database.setVarbinds.append(varbind)
            self.database.setHandlers.append(writer)
            self.database.setUndoData.append(undoData)
        if error != ax.ERR_NOERROR:
            resp = ax.ResponsePDU(True, self.sessionID, packet.transactionID,
                                  packet.packetID, 0, error, bindIndex)
            self.sendPacket(resp, False)
            for i in range(bindIndex):
                # Errored out, clear the successful ones
                self.database.setHandlers[i]("clear",
                                             self.database.setVarbinds[i])
            self.database.inSetP = False
        else:
            resp = ax.ResponsePDU(True, self.sessionID, packet.transactionID,
                                  packet.packetID, 0, ax.ERR_NOERROR, 0)
            self.sendPacket(resp, False)

    def handle_CommitSetPDU(self, packet):
        if not self.database.inSetP:
            pass  # how to handle this?
        varbinds = self.database.setVarbinds
        handlers = self.database.setHandlers
        for i in range(len(varbinds)):
            error = handlers[i]("commit", varbinds[i])
            if error != ax.ERR_NOERROR:
                break
        if error != ax.ERR_NOERROR:
            resp = ax.ResponsePDU(True, self.sessionID, packet.transactionID,
                                  packet.packetID, 0, error, i)
        else:
            resp = ax.ResponsePDU(True, self.sessionID, packet.transactionID,
                                  packet.packetID, 0, ax.ERR_NOERROR, 0)
        self.sendPacket(resp, False)

    def handle_UndoSetPDU(self, packet):
        varbinds = self.database.setVarbinds
        handlers = self.database.setHandlers
        undoData = self.database.setUndoData
        for i in range(len(varbinds)):
            error = handlers[i]("undo", varbinds[i], undoData[i])
            if error != ax.ERR_NOERROR:
                break
        if error != ax.ERR_NOERROR:
            resp = ax.ResponsePDU(True, self.sessionID, packet.transactionID,
                                  packet.packetID, 0, error, i)
        else:
            resp = ax.ResponsePDU(True, self.sessionID, packet.transactionID,
                                  packet.packetID, 0, ax.ERR_NOERROR, 0)
        self.sendPacket(resp, False)

    def handle_CleanupSetPDU(self, packet):
        varbinds = self.database.setVarbinds
        handlers = self.database.setHandlers
        for i in range(len(varbinds)):
            handlers[i]("clean", varbinds[i])
        self.database.inSetP = False

    def handle_ResponsePDU(self, packet):
        index = (packet.sessionID, packet.transactionID, packet.packetID)
        if index in self.packetLog:
            timeout, originalPkt, callback = self.packetLog[index]
            del self.packetLog[index]
            if callback is not None:
                callback(packet, originalPkt)
        else:
            # Ok, response with no associated packet.
            # Probably something that timed out.
            pass


def walkMIBTree(tree, rootpath=()):
    # Tree node formats:
    # {"reader": <func>, "writer": <func>, "subids": {.blah.}}
    # {"reader": <func>, "writer": <func>, "subids": <func>}
    # The "subids" function in dynamic nodes must return an MIB tree
    nodeStack = []
    oidStack = []
    current = tree
    currentKeys = list(current.keys())
    currentKeys.sort()
    keyID = 0
    while True:
        if keyID >= len(currentKeys):
            if nodeStack:
                # No more nodes this level, pop higher node
                current, currentKeys, keyID, key = nodeStack.pop()
                oidStack.pop()
                keyID += 1
                continue
            else:  # Out of tree, we are done
                return
        key = currentKeys[keyID]
        oid = ax.OID(rootpath + tuple(oidStack) + (key,))
        yield (oid, current[key].get("reader"), current[key].get("writer"))
        subs = current[key].get("subids")
        if subs is not None:
            # Push current node, move down a level
            nodeStack.append((current, currentKeys, keyID, key))
            oidStack.append(key)
            if isinstance(subs, dict):
                current = subs
            else:
                current = subs()  # Tree generator function
                if current == {}:  # no dynamic subids, pop
                    current, currentKeys, keyID, key = nodeStack.pop()
                    oidStack.pop()
                    keyID += 1
                    continue
            currentKeys = list(current.keys())
            currentKeys.sort()
            keyID = 0
            key = currentKeys[keyID]
            continue
        keyID += 1
