Python. Two-Factor Authentication with Google Authenticator

"""
Time-based One-time Password Algorithm

Based on the pyotp: https://github.com/nathforge/pyotp
"""
import base64
import hashlib
import hmac
import datetime
import random
import time
import sys
import urllib


def random_base32(length=16, random=random.SystemRandom(),
                  chars=base64._b32alphabet.values()):
    return ''.join(
        random.choice(chars)
        for i in xrange(length)
    )


class OTP(object):
    def __init__(self, secret, digits=6):
        self.secret = secret

    def int_to_bytestring(self, int, padding=8):
        """
        Turns an integer to the OATH specified
        bytestring, which is fed to the HMAC
        along with the secret
        """
        result = []
        while int != 0:
            result.append(chr(int & 0xFF))
            int = int >> 8
        return ''.join(reversed(result)).rjust(padding, '\0')

    def generate_otp(self, input):
        """
        @param [Integer] input the number used seed the HMAC
        Usually either the counter, or the computed integer
        based on the Unix timestamp
        """
        hmac_hash = hmac.new(
            base64.b32decode(self.secret, casefold=True),
            self.int_to_bytestring(input),
            hashlib.sha1,
        ).digest()

        offset = ord(hmac_hash[19]) & 0xf
        code = ((ord(hmac_hash[offset]) & 0x7f) << 24 |
                (ord(hmac_hash[offset + 1]) & 0xff) << 16 |
                (ord(hmac_hash[offset + 2]) & 0xff) << 8 |
                (ord(hmac_hash[offset + 3]) & 0xff))
        # '6' is number of integers in the OTP
        return code % 10 ** 6


class TOTP(OTP):
    def __init__(self, *args, **kwargs):
        """
        @option options [Integer] internval (30) the interval in seconds
            This defaults to 30 which is standard.
        """
        self.interval = kwargs.pop('interval', 30)
        super(TOTP, self).__init__(*args, **kwargs)

    def timecode(self, for_time):
        i = time.mktime(for_time.timetuple())
        return int(i / self.interval)

    def now(self):
        """
        Generate the current time OTP
        @return [Integer] the OTP as an integer
        """
        return self.generate_otp(self.timecode(datetime.datetime.now()))

    def at(self, date):
        """
        Generate the current time OTP
        @return [Integer] the OTP as an integer
        """
        return self.generate_otp(self.timecode(date))

    def verify(self, otp, for_time=None):
        """
        Verifies the OTP passed in against the current time OTP
        @param [String/Integer] otp the OTP to check against
        """
        if for_time is None:
            for_time = datetime.datetime.now()
        return otp == self.at(for_time)
    
    def provisioning_uri(self, name):
        """
        Returns the provisioning URI for the OTP
        This can then be encoded in a QR Code and used
        to provision the Google Authenticator app
        @param [String] name of the account
        @return [String] provisioning uri
        """
        return 'otpauth://totp/%(name)s?secret=%(secret)s' % {
            'name': urllib.quote(name, safe='@'),
            'secret': self.secret,
        }


if __name__ == "__main__":
    secret = "AAAAAAAAAAAAAAAA"
    unixtime = 0
    if len(sys.argv) > 1:
        unixtime = int(sys.argv[1])
    if unixtime > 1:
        date = datetime.datetime.fromtimestamp(unixtime)
    else:
        date = datetime.datetime.now()
    totp = TOTP(secret)
    print "TOTP token for secret '%s' at '%s' is: %s" % (
        secret, date, totp.at(date))

import totp
import qrcode
import datetime


class TotpAuth(object):

    def __init__(self, secret=None):
        if secret is None:
            secret = totp.random_base32()
        self.secret = secret
        self.totp = totp.TOTP(secret)

    def generate_token(self):
        return self.totp.now()

    def valid(self, token):
        try:
            token = int(token)
        except ValueError:
            return False
        now = datetime.datetime.now()
        time30secsago = now + datetime.timedelta(seconds=-30)
        try:
            valid_now = self.totp.verify(token)
            valid_past = self.totp.verify(token, for_time=time30secsago)
            return valid_now or valid_past
        except:
            return False

    def qrcode(self, username):
        uri = self.totp.provisioning_uri(username)
        return qrcode.make(uri, box_size=4)

from totp_auth import TotpAuth

def main(secret=None):
    if not secret:
        # WHEN SIGN UP
        ta = TotpAuth()
    else:
        # WHEN SIGN IN
        ta = TotpAuth(secret)

    # SHOW SECRET AND QR CODE
    qr = ta.qrcode('adw0rd@adw0rd.com')
    print 'SECRET: ', ta.secret
    qr._img.show()

    # CHECK TOKEN
    while True:
        token = raw_input('Please input the token from Google Authenticator: ')
        if ta.valid(token):
            print "SUCCESS"
        else:
            print "FAILED"

if __name__ == '__main__':
    import sys
    main(*sys.argv[1:])