#!/usr/local/bin/python3
#
# Google Authenticator verification script for FreeRADIUS
# Part of pfSense (https://www.pfsense.org)
#
# SPDX-License-Identifier: Apache-2.0
#
# This script validates TOTP codes for FreeRADIUS authentication
# Enhanced with configurable anti-replay protection (RFC 6238 compliance)
#
# Usage: googleauth.py <username> <secret> <pin> <password> [antireplay]
#   antireplay: '1' to enable, '0' or omitted to disable
#

import sys
import time
import struct
import hmac
import hashlib
import base64
import syslog
import os
import fcntl

# Anti-replay configuration
USED_TOKENS_DIR = "/var/run/freeradius"
USED_TOKENS_FILE = os.path.join(USED_TOKENS_DIR, "totp_used_tokens")
TOKEN_EXPIRY_SECONDS = 90  # Tokens expire after 90 seconds (3 TOTP windows)


def ensure_token_dir():
    """Ensure the token directory exists with proper permissions."""
    if not os.path.exists(USED_TOKENS_DIR):
        try:
            os.makedirs(USED_TOKENS_DIR, mode=0o700)
        except OSError:
            pass


def is_token_used(username, token_code):
    """
    Check if the token has already been used (anti-replay protection).
    
    Args:
        username: The username to check
        token_code: The TOTP code to verify
        
    Returns:
        True if token was already used, False otherwise
    """
    ensure_token_dir()
    try:
        if not os.path.exists(USED_TOKENS_FILE):
            return False
        current_time = time.time()
        with open(USED_TOKENS_FILE, 'r') as f:
            for line in f:
                parts = line.strip().split(':')
                if len(parts) == 3:
                    stored_user, stored_token, stored_time = parts
                    if stored_user == username and stored_token == token_code:
                        if current_time - float(stored_time) < TOKEN_EXPIRY_SECONDS:
                            return True
        return False
    except (IOError, ValueError):
        return False


def mark_token_used(username, token_code):
    """
    Record the token as used to prevent replay attacks.
    
    Args:
        username: The username that used the token
        token_code: The TOTP code that was used
    """
    ensure_token_dir()
    try:
        with open(USED_TOKENS_FILE, 'a') as f:
            fcntl.flock(f, fcntl.LOCK_EX)
            f.write("{}:{}:{}\n".format(username, token_code, time.time()))
            fcntl.flock(f, fcntl.LOCK_UN)
        cleanup_expired_tokens()
    except IOError:
        pass


def cleanup_expired_tokens():
    """Remove expired token entries to prevent file growth."""
    try:
        if not os.path.exists(USED_TOKENS_FILE):
            return
        current_time = time.time()
        valid_lines = []
        with open(USED_TOKENS_FILE, 'r') as f:
            for line in f:
                parts = line.strip().split(':')
                if len(parts) == 3:
                    try:
                        if current_time - float(parts[2]) < TOKEN_EXPIRY_SECONDS:
                            valid_lines.append(line)
                    except ValueError:
                        pass
        with open(USED_TOKENS_FILE, 'w') as f:
            fcntl.flock(f, fcntl.LOCK_EX)
            f.writelines(valid_lines)
            fcntl.flock(f, fcntl.LOCK_UN)
    except IOError:
        pass


def authenticate(username, secretkey, pin, password, antireplay_enabled=False):
    """
    Authenticate user with PIN + TOTP code.
    
    Args:
        username: The username attempting authentication
        secretkey: Base32-encoded TOTP secret
        pin: User's static PIN
        password: Full password (PIN + TOTP code concatenated)
        antireplay_enabled: Whether to check for replay attacks
    
    Returns:
        True if authentication successful, False otherwise
    """
    # Validate password length
    if len(password) < len(pin):
        syslog.syslog(syslog.LOG_ERR,
            "freeRADIUS: Google Authenticator - Authentication failed. "
            "User: " + username + ", Reason: wrong PIN")
        return False

    # Verify PIN prefix
    if pin != password[:len(pin)]:
        syslog.syslog(syslog.LOG_ERR,
            "freeRADIUS: Google Authenticator - Authentication failed. "
            "User: " + username + ", Reason: wrong PIN")
        return False

    # Extract the OTP code
    otp_attempt = password[len(pin):]
    
    # Anti-replay check (only if enabled)
    if antireplay_enabled:
        if is_token_used(username, otp_attempt):
            syslog.syslog(syslog.LOG_WARNING,
                "freeRADIUS: Google Authenticator - REPLAY ATTACK DETECTED. "
                "User: " + username + ", Token already used")
            return False

    # Decode the secret key
    try:
        key = base64.b32decode(secretkey)
    except Exception:
        syslog.syslog(syslog.LOG_ERR,
            "freeRADIUS: Google Authenticator - Authentication failed. "
            "User: " + username + ", Reason: invalid secret key")
        return False

    # Get current time window
    tm = int(time.time() / 30)

    # Check current time window and adjacent windows for clock drift tolerance
    for offset in [-1, 0, 1]:
        try:
            b = struct.pack(">q", tm + offset)
            hm = hmac.HMAC(key, b, hashlib.sha1).digest()
            extract_offset = hm[-1] & 0x0F
            truncated = hm[extract_offset:extract_offset + 4]
            code = struct.unpack(">L", truncated)[0]
            code &= 0x7FFFFFFF
            code %= 1000000

            if ("%06d" % code) == str(otp_attempt):
                # Mark token as used (only if anti-replay enabled)
                if antireplay_enabled:
                    mark_token_used(username, otp_attempt)
                syslog.syslog(syslog.LOG_NOTICE,
                    "freeRADIUS: Google Authenticator - "
                    "Authentication successful for user: " + username)
                return True
        except Exception:
            continue

    syslog.syslog(syslog.LOG_ERR,
        "freeRADIUS: Google Authenticator - Authentication failed. "
        "User: " + username + ", Reason: wrong OTP code")
    return False


if __name__ == "__main__":
    # Validate command line arguments
    # Usage: googleauth.py <username> <secret> <pin> <password> [antireplay]
    if len(sys.argv) < 5:
        syslog.syslog(syslog.LOG_ERR,
            "freeRADIUS: Google Authenticator - Invalid number of parameters. "
            "Expected: username secretkey pin password [antireplay]")
        sys.exit(1)

    username = sys.argv[1]
    secretkey = sys.argv[2]
    pin = sys.argv[3]
    password = sys.argv[4]
    
    # Check for anti-replay parameter (5th argument)
    # '1' = enabled, '0' or missing = disabled
    antireplay_enabled = False
    if len(sys.argv) >= 6:
        antireplay_enabled = (sys.argv[5] == '1')

    if authenticate(username, secretkey, pin, password, antireplay_enabled):
        sys.exit(0)
    else:
        sys.exit(1)