#!/usr/bin/python
# -*- Mode: Python -*-
# vi:si:et:sw=4:sts=4:ts=4

import crypt

from twisted.cred import checkers, credentials, error
from twisted.cred.portal import IRealm

from twisted.spread import pb
from twisted.internet import reactor, defer
from twisted.python import failure

class IUsernameCryptPassword(credentials.ICredentials):
    """
    I encapsulate a username and check crypted passwords.

    This credential interface is used when a crypt password is received
    from the party requesting authentication.
    CredentialCheckers which check this kind of credential must store
    the passwords in plaintext or crypt form.

    @type username: C{str}
    @ivar username: The username associated with these credentials.
    """

    def checkPassword(self, password):
        """
        Validate these credentials against the correct plaintext password.
                                                                                
        @param password: The correct, plaintext password against which to
        check.
                                                                                
        @return: a deferred which becomes, or a boolean indicating if the
        password matches.
        """
    
    def checkCryptPassword(self, cryptPassword):
        """
        Validate these credentials against the correct crypt password.
                                                                                
        @param cryptPassword: The correct, crypt password against which to
        check.
                                                                                
        @return: a deferred which becomes, or a boolean indicating if the
        password matches.
        """

class UsernameCryptPasswordPlaintext:
    """
    I take a username and a plaintext password.
    I implement IUsernameCryptPassword.
    """
    
    __implements__ = IUsernameCryptPassword
    def __init__(self, username, password):
        self.username = username
        self.password = password

    def checkPassword(self, password):
        """Check credentials against the given plaintext password."""
        return self.password == password

    def checkCryptPassword(self, cryptPassword):
        """Check credentials against the given cryptPassword."""
        salt = cryptPassword[:2]
        encrypted = crypt.crypt(self.password, salt)
        return encrypted == cryptPassword

class UsernameCryptPasswordCrypt:
    """
    I take a username and a crypt password.
    When using me you should make sure the password was crypted with the
    correct salt (which is stored in the crypt password backend of whatever
    checker you use); otherwise your password may be a valid crypt, but
    with a different salt.
    I implement IUsernameCryptPassword.
    """
    
    __implements__ = IUsernameCryptPassword
    def __init__(self, username, cryptPassword):
        self.username = username
        self.cryptPassword = cryptPassword

    def checkPassword(self, password):
        """Check credentials against the given plaintext password."""
        salt = self.cryptPassword[:2]
        encrypted = crypt.crypt(password, salt)
        return self.cryptPassword == encrypted

    def checkCryptPassword(self, cryptPassword):
        """Check credentials against the given cryptPassword."""
        return self.cryptPassword == cryptPassword

class CryptChecker:
    __implements__ = checkers.ICredentialsChecker
    credentialInterfaces = ( IUsernameCryptPassword, )

    def __init__(self, **users):
        self.users = users

    def addUser(self, username, cryptPassword):
        self.users[username] = cryptPassword

    def _cbCryptPasswordMatch(self, matched, username):
        if matched:
            return username
        else:
            return failure.Failure(error.UnauthorizedLogin())

    def requestAvatarId(self, credentials):
        if credentials.username in self.users:
            return defer.maybeDeferred(
                credentials.checkCryptPassword,
                self.users[credentials.username]).addCallback(
                self._cbCryptPasswordMatch, credentials.username)
        else:
            return failure.Failure(error.UnauthorizedLogin())

def check1(checker):
    # this will succeed
    cred = UsernameCryptPasswordPlaintext('test', 'test')
    d = checker.requestAvatarId(cred)
    d.addCallback(check1Callback, checker)
    d.addErrback(check1Errback)

def check1Callback(avatarId, checker):
    if not avatarId == 'test':
        raise "first password refused, not expected"

    print "first password accepted as expected"

    # next step
    reactor.callLater(0, check2, checker)

def check1Errback(failure):
    print "should not have been reached: %r" % failure

def check2(checker):
    # this will fail
    cred = UsernameCryptPasswordPlaintext('test', 'tes')
    d = checker.requestAvatarId(cred)
    d.addCallback(check2Callback)
    d.addErrback(check2Errback, checker)
    
def check2Callback():
    raise "this one should have failed"

def check2Errback(failure, checker):
    failure.trap(error.UnauthorizedLogin)
    # we got a failed login, which is what we wanted.
    print "second password refused as expected"

    # next step
    reactor.callLater(0, check3, checker)

def check3(checker):
    # this will succeed
    # in a real application you need to figure out the salt, then
    # encrypt your plaintext password you're using
    cred = UsernameCryptPasswordCrypt('test', 'qi1Lftt0GZC0o')
    d = checker.requestAvatarId(cred)
    d.addCallback(check3Callback, checker)
    d.addErrback(check3Errback)

def check3Callback(avatarId, checker):
    if not avatarId == 'test':
        raise "third password refused, not expected"

    print "third password accepted as expected"

    # next step
    reactor.callLater(0, check4, checker)

def check3Errback(failure):
    print "should not have been reached: %r" % failure

def check4(checker):
    # this will fail
    cred = UsernameCryptPasswordCrypt('test', 'qawhatever')
    d = checker.requestAvatarId(cred)
    d.addCallbacks(check4Callback, check4Errback)
    
def check4Callback():
    raise "this one should have failed"

def check4Errback(failure):
    failure.trap(error.UnauthorizedLogin)
    # we got a failed login, which is what we wanted.
    print "fourth password refused as expected"

    # next step
    reactor.callLater(0, reactor.stop)

def main():
    checker = CryptChecker()
    checker.addUser("test", "qi1Lftt0GZC0o")

    reactor.callLater(0, check1, checker)
    reactor.run()

if __name__ == '__main__':
    main()