# Author: Nathaniel McCallum <npmccallum@redhat.com>
#
# Copyright (c) 2013 Red Hat, Inc.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


#
# This script tests OTP, both UDP and Unix Sockets, with a variety of
# configuration. It requires pyrad to run, but exits gracefully if not found.
# It also deliberately shuts down the test daemons between tests in order to
# test how OTP handles the case of short daemon restarts.
#

from k5test import *
from queue import Empty
import io
import struct

try:
    from pyrad import packet, dictionary
except ImportError:
    skip_rest('OTP tests', 'Python pyrad module not found')
try:
    from multiprocessing import Process, Queue
except ImportError:
    skip_rest('OTP tests', 'Python version 2.6 required')

# We could use a dictionary file, but since we need so few attributes,
# we'll just include them here.
radius_attributes = '''
ATTRIBUTE    User-Name    1    string
ATTRIBUTE    User-Password   2    octets
ATTRIBUTE    Service-Type    6    integer
ATTRIBUTE    NAS-Identifier  32    string
'''

class RadiusDaemon(Process):
    MAX_PACKET_SIZE = 4096
    DICTIONARY = dictionary.Dictionary(io.StringIO(radius_attributes))

    def listen(self, addr):
        raise NotImplementedError()

    def recvRequest(self, data):
        raise NotImplementedError()

    def run(self):
        addr = self._args[0]
        secrfile = self._args[1]
        pswd = self._args[2]
        outq = self._args[3]

        if secrfile:
            with open(secrfile, 'rb') as file:
                secr = file.read().strip()
        else:
            secr = b''

        data = self.listen(addr)
        outq.put("started")
        (buf, sock, addr) = self.recvRequest(data)
        pkt = packet.AuthPacket(secret=secr,
                                dict=RadiusDaemon.DICTIONARY,
                                packet=buf)

        usernm = []
        passwd = []
        for key in pkt.keys():
            if key == 'User-Password':
                passwd = list(map(pkt.PwDecrypt, pkt[key]))
            elif key == 'User-Name':
                usernm = pkt[key]

        reply = pkt.CreateReply()
        replyq = {'user': usernm, 'pass': passwd}
        if passwd == [pswd]:
            reply.code = packet.AccessAccept
            replyq['reply'] = True
        else:
            reply.code = packet.AccessReject
            replyq['reply'] = False

        outq.put(replyq)
        if addr is None:
            sock.send(reply.ReplyPacket())
        else:
            sock.sendto(reply.ReplyPacket(), addr)
        sock.close()

class UDPRadiusDaemon(RadiusDaemon):
    def listen(self, addr):
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.bind((addr.split(':')[0], int(addr.split(':')[1])))
        return sock

    def recvRequest(self, sock):
        (buf, addr) = sock.recvfrom(RadiusDaemon.MAX_PACKET_SIZE)
        return (buf, sock, addr)

class UnixRadiusDaemon(RadiusDaemon):
    def listen(self, addr):
        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        if os.path.exists(addr):
            os.remove(addr)
        sock.bind(addr)
        sock.listen(1)
        return (sock, addr)

    def recvRequest(self, sock_and_addr):
        sock, addr = sock_and_addr
        conn = sock.accept()[0]
        sock.close()
        os.remove(addr)

        buf = b''
        remain = RadiusDaemon.MAX_PACKET_SIZE
        while True:
            buf += conn.recv(remain)
            remain = RadiusDaemon.MAX_PACKET_SIZE - len(buf)
            if (len(buf) >= 4):
                remain = struct.unpack("!BBH", buf[0:4])[2] - len(buf)
                if (remain <= 0):
                    return (buf, conn, None)

def verify(daemon, queue, reply, usernm, passwd):
    try:
        data = queue.get(timeout=1)
    except Empty:
        sys.stderr.write("ERROR: Packet not received by daemon!\n")
        daemon.terminate()
        sys.exit(1)
    assert data['reply'] is reply
    assert data['user'] == [usernm]
    assert data['pass'] == [passwd]
    daemon.join()

# Compose a single token configuration.
def otpconfig_1(toktype, username=None, indicators=None):
    val = '{"type": "%s"' % toktype
    if username is not None:
        val += ', "username": "%s"' % username
    if indicators is not None:
        qind = ['"%s"' % s for s in indicators]
        jsonlist = '[' + ', '.join(qind) + ']'
        val += ', "indicators":' + jsonlist
    val += '}'
    return val

# Compose a token configuration list suitable for the "otp" string
# attribute.
def otpconfig(toktype, username=None, indicators=None):
    return '[' + otpconfig_1(toktype, username, indicators) + ']'

prefix = "/tmp/%d" % os.getpid()
secret_file = prefix + ".secret"
socket_file = prefix + ".socket"
with open(secret_file, "w") as file:
    file.write("otptest")
atexit.register(lambda: os.remove(secret_file))

conf = {'plugins': {'kdcpreauth': {'enable_only': 'otp'}},
        'otp': {'udp': {'server': '127.0.0.1:$port9',
                        'secret': secret_file,
                        'strip_realm': 'true',
                        'indicator': ['indotp1', 'indotp2']},
                'unix': {'server': socket_file,
                         'strip_realm': 'false'}}}

queue = Queue()

realm = K5Realm(kdc_conf=conf)
realm.run([kadminl, 'modprinc', '+requires_preauth', realm.user_princ])
flags = ['-T', realm.ccache]
server_addr = '127.0.0.1:' + str(realm.portbase + 9)

## Test UDP fail / custom username
mark('UDP fail / custom username')
daemon = UDPRadiusDaemon(args=(server_addr, secret_file, 'accept', queue))
daemon.start()
queue.get()
realm.run([kadminl, 'setstr', realm.user_princ, 'otp',
           otpconfig('udp', 'custom')])
realm.kinit(realm.user_princ, 'reject', flags=flags, expected_code=1)
verify(daemon, queue, False, 'custom', 'reject')

## Test UDP success / standard username
mark('UDP success / standard username')
daemon = UDPRadiusDaemon(args=(server_addr, secret_file, 'accept', queue))
daemon.start()
queue.get()
realm.run([kadminl, 'setstr', realm.user_princ, 'otp', otpconfig('udp')])
realm.kinit(realm.user_princ, 'accept', flags=flags)
verify(daemon, queue, True, realm.user_princ.split('@')[0], 'accept')
realm.extract_keytab(realm.krbtgt_princ, realm.keytab)
realm.run(['./adata', realm.krbtgt_princ],
          expected_msg='+97: [indotp1, indotp2]')

# Repeat with an indicators override in the string attribute.
mark('auth indicator override')
daemon = UDPRadiusDaemon(args=(server_addr, secret_file, 'accept', queue))
daemon.start()
queue.get()
oconf = otpconfig('udp', indicators=['indtok1', 'indtok2'])
realm.run([kadminl, 'setstr', realm.user_princ, 'otp', oconf])
realm.kinit(realm.user_princ, 'accept', flags=flags)
verify(daemon, queue, True, realm.user_princ.split('@')[0], 'accept')
realm.extract_keytab(realm.krbtgt_princ, realm.keytab)
realm.run(['./adata', realm.krbtgt_princ],
          expected_msg='+97: [indtok1, indtok2]')

# Detect upstream pyrad bug
#   https://github.com/wichert/pyrad/pull/18
try:
    auth = packet.Packet.CreateAuthenticator()
    packet.Packet(authenticator=auth, secret=b'').ReplyPacket()
except AssertionError:
    skip_rest('OTP UNIX domain socket tests', 'pyrad assertion bug detected')

## Test Unix fail / custom username
mark('Unix socket fail / custom username')
daemon = UnixRadiusDaemon(args=(socket_file, None, 'accept', queue))
daemon.start()
queue.get()
realm.run([kadminl, 'setstr', realm.user_princ, 'otp',
           otpconfig('unix', 'custom')])
realm.kinit(realm.user_princ, 'reject', flags=flags, expected_code=1)
verify(daemon, queue, False, 'custom', 'reject')

## Test Unix success / standard username
mark('Unix socket success / standard username')
daemon = UnixRadiusDaemon(args=(socket_file, None, 'accept', queue))
daemon.start()
queue.get()
realm.run([kadminl, 'setstr', realm.user_princ, 'otp', otpconfig('unix')])
realm.kinit(realm.user_princ, 'accept', flags=flags)
verify(daemon, queue, True, realm.user_princ, 'accept')

## Regression test for #8708: test with the standard username and two
## tokens configured, with the first rejecting and the second
## accepting.  With the bug, the KDC incorrectly rejects the request
## and then performs invalid memory accesses, most likely crashing.
daemon1 = UDPRadiusDaemon(args=(server_addr, secret_file, 'accept1', queue))
daemon2 = UnixRadiusDaemon(args=(socket_file, None, 'accept2', queue))
daemon1.start()
queue.get()
daemon2.start()
queue.get()
oconf = '[' + otpconfig_1('udp') + ', ' + otpconfig_1('unix') + ']'
realm.run([kadminl, 'setstr', realm.user_princ, 'otp', oconf])
realm.kinit(realm.user_princ, 'accept2', flags=flags)
verify(daemon1, queue, False, realm.user_princ.split('@')[0], 'accept2')
verify(daemon2, queue, True, realm.user_princ, 'accept2')

success('OTP tests')