Python pyotp.TOTP Examples

def do_mfa_view():

    if 'username' not in g.session:
        return redirect('/user/login')

    if libmfa.mfa_is_enabled(g.session['username']):
        return render_template('mfa.disable.html')
    else:
        libmfa.mfa_reset_secret(g.session['username'])
        secret = libmfa.mfa_get_secret(g.session['username'])
        secret_url = pyotp.totp.TOTP(secret).provisioning_uri(g.session['username'], issuer_name="Vulpy")
        img = qrcode.make(secret_url)

        buffered = BytesIO()
        img.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()

        return render_template('mfa.enable.html', secret_url=secret_url, img_str=img_str) 

def test_auth_mfa_success(flask_client):
    user = User.create(
        email="a@b.c",
        password="password",
        name="Test User",
        activated=True,
        enable_otp=True,
        otp_secret="base32secret3232",
    )
    db.session.commit()

    totp = pyotp.TOTP(user.otp_secret)
    s = Signer(FLASK_SECRET)
    mfa_key = s.sign(str(user.id))

    r = flask_client.post(
        url_for("api.auth_mfa"),
        json={"mfa_token": totp.now(), "mfa_key": mfa_key, "device": "Test Device"},
    )

    assert r.status_code == 200
    assert r.json["api_key"]
    assert r.json["email"]
    assert r.json["name"] == "Test User" 

def test_auth_wrong_mfa_key(flask_client):
    user = User.create(
        email="a@b.c",
        password="password",
        name="Test User",
        activated=True,
        enable_otp=True,
        otp_secret="base32secret3232",
    )
    db.session.commit()

    totp = pyotp.TOTP(user.otp_secret)

    r = flask_client.post(
        url_for("api.auth_mfa"),
        json={
            "mfa_token": totp.now(),
            "mfa_key": "wrong mfa key",
            "device": "Test Device",
        },
    )

    assert r.status_code == 400
    assert r.json["error"] 

def test_request_tfa_token(test_client, authed_sempo_admin_user, otp_generator, status_code):
    """
    GIVEN a Flask Application
    WHEN '/api/auth/tfa/' is requested (POST)
    THEN check a tfa token is only returned when OTP is valid
    """

    auth_token = authed_sempo_admin_user.encode_auth_token().decode()

    tfa_url = authed_sempo_admin_user.tfa_url
    tfa_secret = tfa_url.split("secret=")[1].split('&')[0]
    func = pyotp.TOTP(tfa_secret)
    otp = otp_generator(func)

    otp_expiry_interval = 1
    response = test_client.post('/api/v1/auth/tfa/',
                                headers=dict(Authorization=auth_token, Accept='application/json'),
                                json=dict(
                                    otp=otp,
                                    otp_expiry_interval=otp_expiry_interval
                                ),
                                content_type='application/json', follow_redirects=True)

    assert response.status_code == status_code 

def claim_secret(self, request, pk):
        """Claim and try to get an achievement with its secret."""

        # Get object
        offer = get_object_or_404(self.queryset, id=pk)

        # Check if secret is valid
        secret = request.data['secret']
        if offer.secret and (secret == offer.secret or secret == pyotp.TOTP(offer.secret).now()):
            if request.user.profile.achievements.filter(offer=offer).exists():
                return Response({'message': 'You already have this achievement!'})

            # Create the achievement
            Achievement.objects.create(
                title=offer.title, description=offer.description, admin_note='SECRET',
                body=offer.body, event=offer.event, verified=True, dismissed=True,
                user=request.user.profile, offer=offer)

            return Response({'message': 'Achievement unlocked successfully!'}, 201)

        return forbidden_no_privileges() 

def do_mfa_enable():

    if 'username' not in g.session:
        return redirect('/user/login')

    secret = libmfa.mfa_get_secret(g.session['username'])

    otp = request.form.get('otp')

    totp = pyotp.TOTP(secret)

    if totp.verify(otp):
        libmfa.mfa_enable(g.session['username'])
        return redirect('/mfa/')
    else:
        flash("The OTP was incorrect")
        return redirect('/mfa/')

    return render_template('mfa.enable.html') 

def generate_code(self, secret):
        """
        Generate two-factor authentication code.

        Args:
            secret (str): one time password authentication secret string.

        Returns:
            str: 6 digit two-factor authentication code.

        Examples:
            >>> secret = balena.twofactor_auth.get_otpauth_secret()
            >>> balena.twofactor_auth.generate_code(secret)
            '259975'

        """

        totp = pyotp.TOTP(secret)
        return totp.now() 

def do_mfa_view():

    if 'username' not in g.session:
        return redirect('/user/login')

    if libmfa.mfa_is_enabled(g.session['username']):
        return render_template('mfa.disable.html')
    else:
        libmfa.mfa_reset_secret(g.session['username'])
        secret = libmfa.mfa_get_secret(g.session['username'])
        secret_url = pyotp.totp.TOTP(secret).provisioning_uri(g.session['username'], issuer_name="Vulpy")
        img = qrcode.make(secret_url)

        buffered = BytesIO()
        img.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()

        return render_template('mfa.enable.html', secret_url=secret_url, img_str=img_str) 

def do_mfa_enable():

    if 'username' not in g.session:
        return redirect('/user/login')

    secret = libmfa.mfa_get_secret(g.session['username'])

    otp = request.form.get('otp')

    totp = pyotp.TOTP(secret)

    if totp.verify(otp):
        libmfa.mfa_enable(g.session['username'])
        return redirect('/mfa/')
    else:
        flash("The OTP was incorrect")
        return redirect('/mfa/')

    return render_template('mfa.enable.html') 

def get_qrcode(skey, username, uid):
    filepath = os.path.join(BASE_DIR, 'media/qrcode/{}/'.format(uid))
    if not os.path.exists(filepath):
        os.makedirs(filepath)
    data = pyotp.totp.TOTP(skey).provisioning_uri(username, issuer_name=u'SOMS')
    qr = QRCode(
        version=1,
        error_correction=constants.ERROR_CORRECT_L,
        box_size=6,
        border=4
    )
    try:
        qr.add_data(data)
        qr.make(fit=True)
        img = qr.make_image()
        img.save(filepath + username + '.png')
        return True
    except Exception, e:
        print e
        return False 

def check_totp(self, _, window, entry, new_otp, this_user):
		if not new_otp.verify(entry.get_text().strip()):
			gui_utilities.show_dialog_warning(
				'Incorrect TOTP',
				self.application.get_active_window(),
				'The specified TOTP code is invalid. Make sure your time\n'\
				+ 'is correct, rescan the QR code and try again.'
			)
			return
		self.application.rpc.remote_table_row_set('users', this_user['id'], {'otp_secret': new_otp.secret})
		gui_utilities.show_dialog_info(
			'TOTP Enrollment',
			self.application.get_active_window(),
			'Successfully set the TOTP secret. Your account is now enrolled\n'\
			+ 'in two factor authentication. You will be prompted to enter the\n'
			+ 'value the next time you login.'
		)
		window.destroy() 

def set_otp(self, hash, selector=None, elem=None, otp_type='time', otp_index=1):
        try:
            import pyotp
        except ImportError:
            print("You must install pyotp to use `set_otp`.")
            print("pip install pyotp")
            return

        if not elem:
            elem = self.get_element(selector)

        if otp_type == 'time':
            otp = pyotp.TOTP(hash)
            response = otp.now()
        else:
            otp = pyotp.HOTP(hash)
            response = otp.at(otp_index)

        self.set_value(selector, response, elem=elem) 

def get_mfa(self):
        t = pyotp.TOTP(self.key)
        return t.now() 

def enrollment_setup(self, _):
		rpc = self.application.rpc
		this_user = rpc.graphql_file(user_gql_query, {'name': rpc.username})['db']['user']
		if this_user['otpSecret'] is not None:
			reset = gui_utilities.show_dialog_yes_no(
				'Already Enrolled',
				self.application.get_active_window(),
				'This account is already enrolled in TOTP,\nreset the existing TOTP token?'
			)
			if not reset:
				return
		new_otp = pyotp.TOTP(pyotp.random_base32())
		provisioning_uri = rpc.username + '@' + self.application.config['server'].split(':', 1)[0]
		provisioning_uri = new_otp.provisioning_uri(provisioning_uri) + '&issuer=King%20Phisher'
		bytes_io = io.BytesIO()
		qrcode_ = qrcode.make(provisioning_uri).get_image()
		qrcode_.save(bytes_io, 'PNG')
		pixbuf_loader = GdkPixbuf.PixbufLoader.new()
		pixbuf_loader.write(bytes_io.getvalue())
		pixbuf_loader.close()
		pixbuf = pixbuf_loader.get_pixbuf()

		self.logger.debug('loading gtk builder file from: ' + gtk_builder_file)
		builder = Gtk.Builder()
		builder.add_from_file(gtk_builder_file)
		window = builder.get_object('TOTPEnrollment.window')
		window.set_transient_for(self.application.get_active_window())

		self.application.add_window(window)

		image = builder.get_object('TOTPEnrollment.image_qrcode')
		image.set_from_pixbuf(pixbuf)

		button_check = builder.get_object('TOTPEnrollment.button_check')
		entry_totp = builder.get_object('TOTPEnrollment.entry_totp')
		button_check.connect('clicked', self.check_totp, window, entry_totp, new_otp, this_user)
		entry_totp.connect('activate', self.check_totp, window, entry_totp, new_otp, this_user)

		window.show_all() 

def verify(self, code):
        return pyotp.TOTP(self.secret).verify(code) 

def enrollment_remove(self, _):
		rpc = self.application.rpc
		this_user = rpc.graphql_file(user_gql_query, {'name': rpc.username})['db']['user']
		if this_user['otpSecret'] is None:
			gui_utilities.show_dialog_info(
				'Not Enrolled',
				self.application.get_active_window(),
				'This account is not currently enrolled in two factor\n'\
				+ 'authentication. There are no changes to make.'
			)
			return
		remove = gui_utilities.show_dialog_yes_no(
			'Already Enrolled',
			self.application.get_active_window(),
			'Are you sure you want to unenroll in TOTP? This will remove\n'\
			+ 'two factor authentication on your account.'
		)
		if not remove:
			return
		rpc.remote_table_row_set('users', this_user['id'], {'otp_secret': None})
		gui_utilities.show_dialog_info(
			'TOTP Unenrollment',
			self.application.get_active_window(),
			'Successfully removed the TOTP secret. Your account is now unenrolled\n'\
			+ 'in two factor authentication. You will no longer be prompted to enter\n'\
			+ 'the value when you login.'
		) 

def main():
    arguments = docopt(__doc__, version='generate_code 0.1')

    password = andotp_decrypt.get_password()
    text = None
    if arguments['--old']:
        text = andotp_decrypt.decrypt_aes(password, arguments['ANDOTP_AES_BACKUP_FILE'])
    else:
        text = andotp_decrypt.decrypt_aes_new_format(password, arguments['ANDOTP_AES_BACKUP_FILE'])

    if not text:
        print("Something went wrong while loading %s. Maybe the passphrase was wrong?" % arguments['ANDOTP_AES_BACKUP_FILE'])
        sys.exit(1)
    entries = json.loads(text)

    found = False
    for entry in entries:
        label = entry['label']
        if entry['type'] == 'TOTP':
            if arguments["MATCH_STRING"].lower() in label.lower():
                found = True
                totp = pyotp.TOTP(entry['secret'], interval=entry['period'])
                print("Matched: %s" % label)
                print(totp.now())
                if not arguments["--all"]:
                    # The all flag wasn't provided, i.e. we only wanted one
                    # match, so we can exit.
                    sys.exit(0)
        else:
            print("Unsupported OTP type: %s" % entry["type"])
            sys.exit(2)
    if not found:
        print("No entry matching '%s' found" % arguments["MATCH_STRING"]) 

def initialize(self):
		if not os.access(gtk_builder_file, os.R_OK):
			gui_utilities.show_dialog_error(
				'Plugin Error',
				self.application.get_active_window(),
				"The GTK Builder data file ({0}) is not available.".format(os.path.basename(gtk_builder_file))
			)
			return False
		self.menu_items = {}
		self.add_submenu('Tools > TOTP Self Enrollment')
		self.menu_items['setup'] = self.add_menu_item('Tools > TOTP Self Enrollment > Setup', self.enrollment_setup)
		self.menu_items['remove'] = self.add_menu_item('Tools > TOTP Self Enrollment > Remove', self.enrollment_remove)
		return True 

def main():
    arguments = docopt(__doc__, version='generate_qr_codes 0.1')

    password = andotp_decrypt.get_password()
    text = None
    if arguments['--old']:
        text = andotp_decrypt.decrypt_aes(password, arguments['ANDOTP_AES_BACKUP_FILE'])
    else:
        text = andotp_decrypt.decrypt_aes_new_format(password, arguments['ANDOTP_AES_BACKUP_FILE'])

    if not text:
        print("Something went wrong while loading %s. Maybe the passphrase was wrong?" % arguments['ANDOTP_AES_BACKUP_FILE'])
        sys.exit(1)
    entries = json.loads(text)
    for entry in entries:
        url = None
        issuer = None
        label = entry['label']
        if " - " in label:
            issuer, label = label.split(" - ", 1)
        if entry['type'] == 'TOTP':
            totp = pyotp.TOTP(entry['secret'], interval=entry['period'])
            url = totp.provisioning_uri(label, issuer_name = issuer)
        elif entry['type'] == 'HOTP':
            totp = pyotp.TOTP(entry['secret'])
            url = totp.provisioning_uri(label, issuer_name = issuer)
        if url:
            img = pyqrcode.create(url)
            save_filename = "".join([c for c in label if c.isalpha() or c.isdigit() or c in "@_-"]).strip() + ".svg"
            img.svg(save_filename, scale=4, background='#fff')
            print("Code saved as: %s" % save_filename) 

def render_qr_code(cls, username, key):
        totp = pyotp.TOTP(key)
        uri = totp.provisioning_uri(username, settings.SITE_NAME)

        qr = qrcode.QRCode(box_size=1)
        qr.add_data(uri)
        qr.make(fit=True)

        image = qr.make_image(fill_color='black', back_color='white')
        buf = BytesIO()
        image.save(buf, format='PNG')
        return 'data:image/png;base64,' + base64.b64encode(buf.getvalue()).decode('ascii') 

def get_mfa(self):
        t = pyotp.TOTP(self.key)
        return t.now() 

def get_mfa(self):
        t = pyotp.TOTP(self.key)
        return t.now() 

def test_totp_claim(self):
        offer_1 = OfferedAchievement.objects.create(
            title="Test Achievement", body=self.body_1, event=self.event_1)
        offer_2 = OfferedAchievement.objects.create(
            title="Test Achievement", body=self.body_1, event=self.event_1)

        # Setup data
        data = {
            'secret': 'something'
        }
        url = '/api/achievements-offer/%s' % offer_1.id

        # Try with invalid secret
        response = self.client.post(url, data, format='json')
        self.assertEqual(response.status_code, 403)

        # Try with master secret
        data['secret'] = offer_1.secret
        response = self.client.post(url, data, format='json')
        self.assertEqual(response.status_code, 201)

        # Try to get again master secret
        response = self.client.post(url, data, format='json')
        self.assertEqual(response.status_code, 200)

        # Try with TOTP for offer 2
        url = '/api/achievements-offer/%s' % offer_2.id
        data['secret'] = pyotp.TOTP(offer_2.secret).now()
        response = self.client.post(url, data, format='json')
        self.assertEqual(response.status_code, 201)
        response = self.client.post(url, data, format='json')
        self.assertEqual(response.status_code, 200) 

def verify_totp(self, token):
        totp = pyotp.TOTP(self.otp_secret)
        return totp.verify(token) 

def soms_mfa(skey, verify_code):
    t = pyotp.TOTP(skey)
    result = t.verify(verify_code)

    return result 

def disable_2fa(user, password, totp_secret, base_url):
    prefix = "/vpn-admin-portal"
    admin_url = base_url + prefix
    browser = mechanicalsoup.StatefulBrowser(raise_on_404=True)
    logger.info("opening auth_url")
    response = browser.open(admin_url)
    assert response.ok
    browser.select_form()
    browser["userName"] = user
    browser["userPass"] = password
    logger.info("logging in")
    response = browser.submit_selected()
    assert response.ok
    form = browser.select_form()
    if form.form.attrs['action'] != prefix + '/_two_factor/auth/verify/totp':
        logger.warning("2fa not enabled")
        return

    # redirected to totp screen
    totp = TOTP(totp_secret)
    browser['_two_factor_auth_totp_key'] = totp.now()
    logger.info("submitting totp key")
    response = browser.submit_selected()
    assert response.ok

    form = browser.select_form()
    if form.form.attrs['action'] == prefix + '/_two_factor/auth/verify/totp':
        error = browser.get_current_page().findAll("p", {"class": "error"})[0].contents[0].strip()
        raise EduvpnAuthException(error)

    response = browser.open("{}/user?user_id={}".format(admin_url, user))
    assert response.ok
    form = browser.select_form()
    button = form.form.select('button[value="deleteTotpSecret"]')
    if button:
        response = browser.submit_selected()
        assert(response.ok)
    else:
        logger.error(form.form)
        logger.error("2fa not enabled, but had to supply otp during login") 

def test_2fa_enroll(self):
        username, password = online_tests
        disable_2fa(username, password, totp_secret=TOTP_SECRET, base_url=INSTANCE_URI)
        oauth, meta = get_oauth_token(username, password, instance_uri=INSTANCE_URI)
        two_factor_enroll_totp(oauth, meta.api_base_uri, secret=TOTP_SECRET, key=TOTP(TOTP_SECRET).now())
        disable_2fa(username, password, totp_secret=TOTP_SECRET, base_url=INSTANCE_URI) 

def check_second_factor(self, second_factor):
    totp = pyotp.TOTP(self.two_factor_secret)
    if not totp.verify(int(second_factor)):
      return False
    return True 

def mfa_cancel():
    if not current_user.enable_otp:
        flash("you don't have MFA enabled", "warning")
        return redirect(url_for("dashboard.index"))

    otp_token_form = OtpTokenForm()
    totp = pyotp.TOTP(current_user.otp_secret)

    if otp_token_form.validate_on_submit():
        token = otp_token_form.token.data

        if totp.verify(token):
            current_user.enable_otp = False
            current_user.otp_secret = None
            db.session.commit()

            # user does not have any 2FA enabled left, delete all recovery codes
            if not current_user.two_factor_authentication_enabled():
                RecoveryCode.empty(current_user)

            flash("MFA is now disabled", "warning")
            return redirect(url_for("dashboard.index"))
        else:
            flash("Incorrect token", "warning")

    return render_template("dashboard/mfa_cancel.html", otp_token_form=otp_token_form) 

def rpc_login(handler, session, username, password, otp=None):
	logger = logging.getLogger('KingPhisher.Server.Authentication')
	if not ipaddress.ip_address(handler.client_address[0]).is_loopback:
		logger.warning("failed login request from {0} for user {1}, (invalid source address)".format(handler.client_address[0], username))
		raise ValueError('invalid source address for login')
	fail_default = (False, ConnectionErrorReason.ERROR_INVALID_CREDENTIALS, None)
	fail_otp = (False, ConnectionErrorReason.ERROR_INVALID_OTP, None)

	if not (username and password):
		logger.warning("failed login request from {0} for user {1}, (missing username or password)".format(handler.client_address[0], username))
		return fail_default
	if not handler.server.forked_authenticator.authenticate(username, password):
		logger.warning("failed login request from {0} for user {1}, (authentication failed)".format(handler.client_address[0], username))
		return fail_default

	user = session.query(db_models.User).filter_by(name=username).first()
	if not user:
		logger.info('creating new user object with name: ' + username)
		user = db_models.User(name=username)
	elif user.has_expired:
		logger.warning("failed login request from {0} for user {1}, (user has expired)".format(handler.client_address[0], username))
		return fail_default
	elif user.otp_secret:
		if otp is None:
			logger.debug("failed login request from {0} for user {1}, (missing otp)".format(handler.client_address[0], username))
			return fail_otp
		if not (isinstance(otp, str) and len(otp) == 6 and otp.isdigit()):
			logger.warning("failed login request from {0} for user {1}, (invalid otp)".format(handler.client_address[0], username))
			return fail_otp
		totp = pyotp.TOTP(user.otp_secret)
		now = datetime.datetime.now()
		if otp not in (totp.at(now + datetime.timedelta(seconds=offset)) for offset in (0, -30, 30)):
			logger.warning("failed login request from {0} for user {1}, (invalid otp)".format(handler.client_address[0], username))
			return fail_otp
	user.last_login = db_models.current_timestamp()
	session.add(user)
	session.commit()
	session_id = handler.server.session_manager.put(user)
	logger.info("successful login request from {0} for user {1} (id: {2})".format(handler.client_address[0], username, user.id))
	signals.send_safe('rpc-user-logged-in', logger, handler, session=session_id, name=username)
	return True, ConnectionErrorReason.SUCCESS, session_id