|
| 1 | +import pyotp |
| 2 | +import argparse |
| 3 | +import json |
| 4 | + |
| 5 | +SECRET_STORAGE_FILE = "secret.json" |
| 6 | + |
| 7 | +def generate_totp(secret_key, algorithm='SHA1', digits=6, interval=30): |
| 8 | + totp = pyotp.TOTP(secret_key, digits=digits, interval=interval, digest=algorithm) |
| 9 | + return totp.now() |
| 10 | + |
| 11 | +def verify_totp(secret_key, code, algorithm='SHA1', digits=6, interval=30, window=1): |
| 12 | + totp = pyotp.TOTP(secret_key, digits=digits, interval=interval, digest=algorithm) |
| 13 | + return totp.verify(code, valid_window=window) |
| 14 | + |
| 15 | +def save_secret(secret_key, filename): |
| 16 | + data = {'secret_key': secret_key} |
| 17 | + with open(filename, 'w') as file: |
| 18 | + json.dump(data, file) |
| 19 | + print(f'Secret key saved to {filename}') |
| 20 | + |
| 21 | +def load_secret(filename): |
| 22 | + try: |
| 23 | + with open(filename, 'r') as file: |
| 24 | + data = json.load(file) |
| 25 | + return data['secret_key'] |
| 26 | + except FileNotFoundError: |
| 27 | + print(f'Secret key file "{filename}" not found.') |
| 28 | + return None |
| 29 | + |
| 30 | +def main(): |
| 31 | + parser = argparse.ArgumentParser(description='Multi-Factor Authentication (MFA) Generator') |
| 32 | + parser.add_argument('--generate', action='store_true', help='Generate a TOTP code') |
| 33 | + parser.add_argument('--verify', help='Verify a TOTP code') |
| 34 | + parser.add_argument('--algorithm', choices=['SHA1', 'SHA256', 'SHA512'], default='SHA1', help='Hash algorithm for TOTP') |
| 35 | + parser.add_argument('--digits', type=int, choices=[6, 8], default=6, help='Number of digits in TOTP code') |
| 36 | + parser.add_argument('--interval', type=int, default=30, help='Time interval for TOTP code generation') |
| 37 | + parser.add_argument('--window', type=int, default=1, help='Verification window for TOTP codes') |
| 38 | + parser.add_argument('--save', help='Save secret key to a file') |
| 39 | + parser.add_argument('--load', action='store_true', help='Load secret key from a file') |
| 40 | + args = parser.parse_args() |
| 41 | + |
| 42 | + if args.load: |
| 43 | + secret_key = load_secret(SECRET_STORAGE_FILE) |
| 44 | + if secret_key is None: |
| 45 | + return |
| 46 | + elif args.generate or args.verify: |
| 47 | + secret_key = input('Enter your secret key: ').strip() |
| 48 | + if args.save: |
| 49 | + save_secret(secret_key, args.save) |
| 50 | + else: |
| 51 | + print('Please specify either --generate or --verify.') |
| 52 | + return |
| 53 | + |
| 54 | + if args.generate: |
| 55 | + code = generate_totp(secret_key, args.algorithm, args.digits, args.interval) |
| 56 | + print(f'Generated TOTP code: {code}') |
| 57 | + |
| 58 | + if args.verify: |
| 59 | + code_to_verify = args.verify |
| 60 | + result = verify_totp(secret_key, code_to_verify, args.algorithm, args.digits, args.interval, args.window) |
| 61 | + if result: |
| 62 | + print('TOTP code is valid.') |
| 63 | + else: |
| 64 | + print('TOTP code is NOT valid.') |
| 65 | + |
| 66 | +if __name__ == '__main__': |
| 67 | + main() |
0 commit comments