ungleich-otp/ungleichotp/otpauth/serializer.py

from rest_framework import serializers, exceptions
from otpauth.models import OTPSeed
import pyotp
import otpauth

class OTPSerializer(serializers.ModelSerializer):
    class Meta:
        model = OTPSeed
        fields = ('name', 'realm', 'seed')
        read_only_fields = ('seed',)

    def create(self, validated_data):
        validated_data['seed'] = pyotp.random_base32()
        return OTPSeed.objects.create(**validated_data)

class VerifySerializerV1(serializers.Serializer):
    """
    This is the first version of the serializer that would authenticate the request
    itself. This is not necessary anymore starting from version 0.5

    Code to be removed prior to 1.0
    """

    name = serializers.CharField(max_length=128)
    token = serializers.CharField(max_length=128)
    realm = serializers.CharField(max_length=128)

    verifyname = serializers.CharField(max_length=128)
    verifytoken = serializers.CharField(max_length=128)
    verifyrealm = serializers.CharField(max_length=128)

    def save(self):
        token_in = self.validated_data.get('token')
        name_in =  self.validated_data.get('name')
        realm_in =  self.validated_data.get('realm')
        verifytoken = self.validated_data.get('verifytoken')
        verifyname =  self.validated_data.get('verifyname')
        verifyrealm =  self.validated_data.get('verifyrealm')

        # 1. Verify that the connection might authenticate
        try:
            db_instance = otpauth.models.OTPSeed.objects.get(name=name_in, realm=realm_in)
        except (OTPSeed.MultipleObjectsReturned, OTPSeed.DoesNotExist):
            raise exceptions.AuthenticationFailed()

        totp = pyotp.TOTP(db_instance.seed)

        if not totp.verify(token_in, valid_window=3):
            raise exceptions.AuthenticationFailed()


        # 2. Verify the requested data
        try:
            verifyinstance = otpauth.models.OTPSeed.objects.get(name=verifyname, realm=verifyrealm)
        except (OTPSeed.MultipleObjectsReturned, OTPSeed.DoesNotExist):
            raise exceptions.PermissionDenied()

        totp = pyotp.TOTP(verifyinstance.seed)

        if not totp.verify(verifytoken, valid_window=3):
            raise exceptions.PermissionDenied()

        print("All verified!")
        return verifyinstance


class TokenSerializer(serializers.Serializer):
    name = serializers.CharField(max_length=128)
    token = serializers.CharField(max_length=128)
    realm = serializers.CharField(max_length=128)

    token_name = 'token'
    name_name = 'name'
    realm_name = 'realm'

    def save(self):
        token_in = self.validated_data.get('token')
        name_in =  self.validated_data.get('name')
        realm_in =  self.validated_data.get('realm')

        # 1. Verify that the connection might authenticate
        try:
            db_instance = otpauth.models.OTPSeed.objects.get(name=name_in, realm=realm_in)
        except (OTPSeed.MultipleObjectsReturned, OTPSeed.DoesNotExist):
            raise exceptions.AuthenticationFailed()

        totp = pyotp.TOTP(db_instance.seed)

        if not totp.verify(token_in, valid_window=3):
            raise exceptions.AuthenticationFailed()

        return (db_instance, token_in)

class VerifySerializer(TokenSerializer):
    token_name = 'verifytoken'
    name_name = 'verifyname'
    realm_name = 'verifyrealm'