Add some rate limiting logic

This commit is contained in:
Quantum 2025-07-20 17:38:02 -04:00
parent c0395eb97a
commit b9d6630b1a
2 changed files with 13 additions and 6 deletions

View file

@ -11,6 +11,8 @@ from typing import Optional
import gssapi import gssapi
import ldap import ldap
from flask import Flask, Response, redirect, request from flask import Flask, Response, redirect, request
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address
from gssapi.exceptions import BadMechanismError, GSSError, GeneralError from gssapi.exceptions import BadMechanismError, GSSError, GeneralError
from ldap.filter import escape_filter_chars from ldap.filter import escape_filter_chars
from werkzeug.routing import Rule from werkzeug.routing import Rule
@ -20,6 +22,10 @@ app.logger.setLevel(logging.INFO)
app.url_map.add(Rule('/krbauth', endpoint='krbauth.auth')) app.url_map.add(Rule('/krbauth', endpoint='krbauth.auth'))
app.url_map.add(Rule('/krbauth/check', endpoint='krbauth.check')) app.url_map.add(Rule('/krbauth/check', endpoint='krbauth.check'))
LIMITER_STORAGE = os.environ.get('KRBAUTH_LIMITER_STORAGE', 'memory://')
LIMITER_FREQUENCY = os.environ.get('KRBAUTH_LIMITER_FREQUENCY', '3/minute')
limiter = Limiter(get_remote_address, app=app, storage_uri=LIMITER_STORAGE)
timestamp = struct.Struct('!q') timestamp = struct.Struct('!q')
hmac_digest = hashlib.sha512 hmac_digest = hashlib.sha512
digest_size = hmac_digest().digest_size digest_size = hmac_digest().digest_size
@ -140,7 +146,7 @@ def auth_spnego(context: Context, next_url: str) -> Response:
) )
result = ldap_ctx.search_s(LDAP_SEARCH_BASE, ldap.SCOPE_SUBTREE, ldap_filter, ['cn']) result = ldap_ctx.search_s(LDAP_SEARCH_BASE, ldap.SCOPE_SUBTREE, ldap_filter, ['cn'])
if not result: if not result:
return make_401('Failed to authenticate', krb5_name=krb5_name) return make_401('Authentication failed', krb5_name=krb5_name)
app.logger.info('Authenticated via Kerberos as: %s, %s', krb5_name, result[0][0]) app.logger.info('Authenticated via Kerberos as: %s, %s', krb5_name, result[0][0])
else: else:
app.logger.info('Authenticated via Kerberos as: %s', krb5_name) app.logger.info('Authenticated via Kerberos as: %s', krb5_name)
@ -164,16 +170,16 @@ def auth_basic(context: Context, next_url: str) -> Response:
try: try:
ldap_ctx.bind_s(dn, password) ldap_ctx.bind_s(dn, password)
except ldap.INVALID_CREDENTIALS: except ldap.INVALID_CREDENTIALS:
return make_401('Failed to authenticate to LDAP', dn=dn) return make_401('Authentication failed', dn=dn)
if context.ldap_group: if context.ldap_group:
if not ldap_ctx.search_s(dn, ldap.SCOPE_BASE, '(memberof=%s)' % ( if not ldap_ctx.search_s(dn, ldap.SCOPE_BASE, '(memberof=%s)' % (
escape_filter_chars(context.ldap_group), escape_filter_chars(context.ldap_group),
)): )):
return make_401('Failed to authenticate to LDAP', dn=dn, group=context.ldap_group) return make_401('Authentication failed', dn=dn, group=context.ldap_group)
app.logger.info('Authenticated via LDAP as: %s in %s', dn, context.ldap_group) app.logger.info('Authenticated via LDAP as: %s in %s from %s', dn, context.ldap_group, request.remote_addr)
else: else:
app.logger.info('Authenticated via LDAP as: %s', dn) app.logger.info('Authenticated via LDAP as: %s from %s', dn, request.remote_addr)
return auth_success(context, next_url) return auth_success(context, next_url)
@ -186,6 +192,7 @@ def check_tls() -> bool:
@app.endpoint('krbauth.auth') @app.endpoint('krbauth.auth')
@limiter.limit(LIMITER_FREQUENCY)
def auth() -> Response: def auth() -> Response:
next_url = request.args.get('next', '/') next_url = request.args.get('next', '/')
context = Context.from_request() context = Context.from_request()

View file

@ -9,7 +9,7 @@ setup(
name='nginx_krbauth', name='nginx_krbauth',
version='0.0.4', version='0.0.4',
py_modules=['nginx_krbauth'], py_modules=['nginx_krbauth'],
install_requires=['flask', 'gssapi', 'python-ldap'], install_requires=['flask', 'gssapi', 'python-ldap', 'flask-limiter'],
author='quantum', author='quantum',
author_email='quantum2048@gmail.com', author_email='quantum2048@gmail.com',