verify_otp_in_message.py

verify_otp_in_message.py

import hmac
import hashlib
import time
import base64
import struct
import sys
import re

# The totp shared secret
shared_secret = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"

def extract_otp_from_message(user_provided_message):
    # Define the regex pattern for 6 digits at the start of the string
    pattern = r"^\d{6}"
    
    # Use re.match to find the match at the beginning of the string
    match = re.match(pattern, user_provided_message)
    
    if match:
        # Extract the 6 digits
        user_provided_totp = match.group()
        
        # Extract the remaining string after the 6 digits
        remaining_string = user_provided_message[len(user_provided_totp):].strip()
        cleaned_remaining_string = re.sub(r"^[;:_\-\s]", "", remaining_string)
        
        return user_provided_totp, cleaned_remaining_string
    else:
        return None, user_provided_message

def TOTP(K, digits=6, timestep=30, num_totps=3):
    """
    Generates a 6-digit TOTP using SHA-1 and a base32 secret key.
    """
    K_bytes = base64.b32decode(K, casefold=True)
    current_time = int(time.time())
    totps = []

    for i in range(num_totps):
        counter = (current_time - (num_totps // 2 - i) * 30) // timestep
        hmac_sha1 = hmac.new(key=K_bytes, msg=struct.pack(">Q", counter), digestmod=hashlib.sha1).digest()
        offset = hmac_sha1[-1] & 0x0F
        truncated_hash = struct.unpack(">I", hmac_sha1[offset:offset + 4])[0] & 0x7FFFFFFF
        otp = str(truncated_hash % 10**digits).zfill(digits)
        totps.append(otp)

    return totps

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python verify_otp_in_message.py <message>")
        print(sys.argv)
        # print(sys.argv[1:])
        sys.exit(2)

    user_provided_message = sys.argv[1]
    user_provided_totp, remaining_string = extract_otp_from_message(user_provided_message)

    if user_provided_totp is not None:
        # Generate the TOTP for the given shared secret
        # print(user_provided_totp)
        generated_totps = TOTP(shared_secret)
        # print(f"Generated TOTP: {generated_totp}")

        if user_provided_totp in generated_totps:
            print("TOTP validation successful! 🎉")
            print(remaining_string)
            sys.exit(0)
        else:
            print("TOTP validation failed. 😢")
            print(remaining_string)
            sys.exit(1)
    else:
        print("NO TOTP provided. 👀")
        print(remaining_string)
        sys.exit(2)