#
# This file is part of pysnmp software.
#
# Copyright (c) 2005-2016, Ilya Etingof <ilya@glas.net>
# License: http://pysnmp.sf.net/license.html
#
import time, sys
from pysnmp.proto.secmod.base import AbstractSecurityModel
from pysnmp.proto.secmod.rfc3414.auth import hmacmd5, hmacsha, noauth
from pysnmp.proto.secmod.rfc3414.priv import des, nopriv
from pysnmp.proto.secmod.rfc3826.priv import aes
from pysnmp.proto.secmod.eso.priv import des3, aes192, aes256
from pysnmp.smi.error import NoSuchInstanceError
from pysnmp.proto import rfc1155, errind, error
from pysnmp import debug
from pyasn1.type import univ, namedtype, constraint
from pyasn1.codec.ber import encoder, decoder, eoo
from pyasn1.error import PyAsn1Error
from pyasn1.compat.octets import null

# USM security params

class UsmSecurityParameters(rfc1155.TypeCoercionHackMixIn, univ.Sequence):
    componentType = namedtype.NamedTypes(
        namedtype.NamedType('msgAuthoritativeEngineId', univ.OctetString()),
        namedtype.NamedType('msgAuthoritativeEngineBoots', univ.Integer().subtype(subtypeSpec=constraint.ValueRangeConstraint(0, 2147483647))),
        namedtype.NamedType('msgAuthoritativeEngineTime', univ.Integer().subtype(subtypeSpec=constraint.ValueRangeConstraint(0, 2147483647))),
        namedtype.NamedType('msgUserName', univ.OctetString().subtype(subtypeSpec=constraint.ValueSizeConstraint(0, 32))),
        namedtype.NamedType('msgAuthenticationParameters', univ.OctetString()),
        namedtype.NamedType('msgPrivacyParameters', univ.OctetString())
    )

class SnmpUSMSecurityModel(AbstractSecurityModel):
    securityModelID = 3
    authServices = {hmacmd5.HmacMd5.serviceID: hmacmd5.HmacMd5(),
                    hmacsha.HmacSha.serviceID: hmacsha.HmacSha(),
                    noauth.NoAuth.serviceID: noauth.NoAuth()}
    privServices = {des.Des.serviceID: des.Des(),
                    des3.Des3.serviceID: des3.Des3(),
                    aes.Aes.serviceID: aes.Aes(),
                    aes192.Aes192.serviceID: aes192.Aes192(),
                    aes256.Aes256.serviceID: aes256.Aes256(),
                    nopriv.NoPriv.serviceID: nopriv.NoPriv()}
    def __init__(self):
        AbstractSecurityModel.__init__(self)
        self.__securityParametersSpec = UsmSecurityParameters()
        self.__timeline = {}
        self.__timelineExpQueue = {}
        self.__expirationTimer = 0
        self.__paramsBranchId = -1

    def __sec2usr(self, snmpEngine, securityName, securityEngineID=None):
        usmUserEngineID, = snmpEngine.msgAndPduDsp.mibInstrumController.mibBuilder.importSymbols('SNMP-USER-BASED-SM-MIB', 'usmUserEngineID')
        if self.__paramsBranchId != usmUserEngineID.branchVersionId:
            usmUserName, usmUserSecurityName = snmpEngine.msgAndPduDsp.mibInstrumController.mibBuilder.importSymbols('SNMP-USER-BASED-SM-MIB', 'usmUserName', 'usmUserSecurityName')

            self.__securityToUserMap = {}

            nextMibNode = usmUserEngineID

            while True:
                try:
                    nextMibNode = usmUserEngineID.getNextNode(nextMibNode.name)

                except NoSuchInstanceError:
                    self.__paramsBranchId = usmUserEngineID.branchVersionId
                    debug.logger & debug.flagSM and debug.logger('_sec2usr: built snmpEngineId + securityName to userName map, version %s: %r' % (self.__paramsBranchId, self.__securityToUserMap))
                    break

                instId = nextMibNode.name[len(usmUserSecurityName.name):]

                __engineID = usmUserEngineID.getNode(usmUserEngineID.name + instId).syntax
                __userName = usmUserName.getNode(usmUserName.name + instId).syntax
                __securityName = usmUserSecurityName.getNode(usmUserSecurityName.name + instId).syntax

                k = __engineID, __securityName

                # first (lesser) securityName wins
                if k not in self.__securityToUserMap:
                    self.__securityToUserMap[k] = __userName

        if securityEngineID is None:
            snmpEngineID, = snmpEngine.msgAndPduDsp.mibInstrumController.mibBuilder.importSymbols('__SNMP-FRAMEWORK-MIB', 'snmpEngineID')
            securityEngineID = snmpEngineID.syntax

        try:
            userName = self.__securityToUserMap[(securityEngineID, securityName)]
        except KeyError:
            debug.logger & debug.flagSM and debug.logger('_sec2usr: no entry exists for snmpEngineId %r, securityName %r' % (securityEngineID, securityName))
            raise NoSuchInstanceError()  # emulate MIB lookup

        debug.logger & debug.flagSM and debug.logger('_sec2usr: using userName %r for snmpEngineId %r, securityName %r' % (userName, securityEngineID, securityName))

        return userName

    def __getUserInfo(self, mibInstrumController, securityEngineID, userName):
        usmUserEntry, = mibInstrumController.mibBuilder.importSymbols(
            'SNMP-USER-BASED-SM-MIB', 'usmUserEntry'
        )
        tblIdx = usmUserEntry.getInstIdFromIndices(securityEngineID, userName)
        # Get userName & securityName
        usmUserName = usmUserEntry.getNode(usmUserEntry.name + (2,) + tblIdx).syntax
        usmUserSecurityName = usmUserEntry.getNode(usmUserEntry.name + (3,) + tblIdx).syntax
        # Get protocols
        usmUserAuthProtocol = usmUserEntry.getNode(usmUserEntry.name + (5,) + tblIdx).syntax
        usmUserPrivProtocol = usmUserEntry.getNode(usmUserEntry.name + (8,) + tblIdx).syntax
        # Get keys
        pysnmpUsmKeyEntry, = mibInstrumController.mibBuilder.importSymbols(
            'PYSNMP-USM-MIB', 'pysnmpUsmKeyEntry'
        )
        pysnmpUsmKeyAuthLocalized = pysnmpUsmKeyEntry.getNode(pysnmpUsmKeyEntry.name + (1,) + tblIdx).syntax
        pysnmpUsmKeyPrivLocalized = pysnmpUsmKeyEntry.getNode(pysnmpUsmKeyEntry.name + (2,) + tblIdx).syntax
        return (usmUserName, usmUserSecurityName, usmUserAuthProtocol,
                pysnmpUsmKeyAuthLocalized, usmUserPrivProtocol,
                pysnmpUsmKeyPrivLocalized)

    def __cloneUserInfo(self, mibInstrumController, securityEngineID,
                        userName):
        snmpEngineID, = mibInstrumController.mibBuilder.importSymbols(
            '__SNMP-FRAMEWORK-MIB', 'snmpEngineID'
        )
        # Proto entry
        usmUserEntry, = mibInstrumController.mibBuilder.importSymbols(
            'SNMP-USER-BASED-SM-MIB', 'usmUserEntry'
        )
        tblIdx1 = usmUserEntry.getInstIdFromIndices(
            snmpEngineID.syntax, userName
        )
        # Get proto protocols
        usmUserName = usmUserEntry.getNode(usmUserEntry.name + (2,) + tblIdx1)
        usmUserSecurityName = usmUserEntry.getNode(usmUserEntry.name + (3,) + tblIdx1)
        usmUserCloneFrom = usmUserEntry.getNode(usmUserEntry.name + (4,) + tblIdx1)
        usmUserAuthProtocol = usmUserEntry.getNode(usmUserEntry.name + (5,) + tblIdx1)
        usmUserPrivProtocol = usmUserEntry.getNode(usmUserEntry.name + (8,) + tblIdx1)
        # Get proto keys
        pysnmpUsmKeyEntry, = mibInstrumController.mibBuilder.importSymbols(
            'PYSNMP-USM-MIB', 'pysnmpUsmKeyEntry'
        )
        pysnmpUsmKeyAuth = pysnmpUsmKeyEntry.getNode(pysnmpUsmKeyEntry.name + (3,) + tblIdx1)
        pysnmpUsmKeyPriv = pysnmpUsmKeyEntry.getNode(pysnmpUsmKeyEntry.name + (4,) + tblIdx1)

        # Create new row from proto values

        tblIdx2 = usmUserEntry.getInstIdFromIndices(securityEngineID, userName)

        # New row
        mibInstrumController.writeVars(
            ((usmUserEntry.name + (13,) + tblIdx2, 4),)
        )

        # Set user&securityNames
        usmUserEntry.getNode(usmUserEntry.name + (2,) + tblIdx2).syntax = usmUserName.syntax
        usmUserEntry.getNode(usmUserEntry.name + (3,) + tblIdx2).syntax = usmUserSecurityName.syntax

        # Store a reference to original row
        usmUserEntry.getNode(usmUserEntry.name + (4,) + tblIdx2).syntax = usmUserCloneFrom.syntax.clone(tblIdx1)

        # Set protocols
        usmUserEntry.getNode(usmUserEntry.name + (5,) + tblIdx2).syntax = usmUserAuthProtocol.syntax
        usmUserEntry.getNode(usmUserEntry.name + (8,) + tblIdx2).syntax = usmUserPrivProtocol.syntax

        # Localize and set keys
        pysnmpUsmKeyEntry, = mibInstrumController.mibBuilder.importSymbols(
            'PYSNMP-USM-MIB', 'pysnmpUsmKeyEntry'
        )
        pysnmpUsmKeyAuthLocalized = pysnmpUsmKeyEntry.getNode(
            pysnmpUsmKeyEntry.name + (1,) + tblIdx2
        )