diff --git a/main.py b/main.py index 5f9209ef3c5d432f6830abf0b71fac72bff9725a..30b05b871d19305d01e6c1c523211df66cbaff7b 100755 --- a/main.py +++ b/main.py @@ -1,15 +1,70 @@ #!/usr/bin/env python3 -import pyotp +import argparse + import keyring +import pyotp + + +def delete_token(): + if keyring.get_password("unihd", "2fa") is None: + return + keyring.delete_password("unihd", "2fa") + def enroll_token(token_uri: str) -> None: keyring.set_password("unihd", "2fa", token_uri) + +def enroll_token_interactive(): + mfa_uri = None + + print("Enrolling secret.") + while mfa_uri is None or not mfa_uri: + mfa_uri = input("Enter/Paste MFA URI: ") + if mfa_uri.startswith("otpauth://"): + enroll_token(mfa_uri) + else: + print("The given string has to start with otpauth://") + + def get_otp() -> str: mfa_uri = keyring.get_password("unihd", "2fa") - otp = pyotp.parse_uri(mfa_uri) - return otp.now() + if mfa_uri is None: + print("No token enrolled, run again with -e or --enroll") + exit(1) + else: + otp = pyotp.parse_uri(mfa_uri) + return otp.now() + if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "-d", + "--delete", + help="Delete the stored otpsecret from keyring", + default=False, + action="store_true", + ) + parser.add_argument( + "-e", + "--enroll", + help="Enrolls a given token to the keyring", + default=None, + const="--", + nargs="?", + ) + args = parser.parse_args() + if args.delete: + delete_token() + exit() + if args.enroll: + if args.enroll == "--": + enroll_token_interactive() + elif args.enroll.startswith("otpauth://"): + enroll_token(args.enroll) + else: + print("The given string has to start with otpauth://") + exit(1) print(get_otp())