diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..f5df813 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,39 @@ +name: Test + +on: + push: + branches: [ master, feature/python-otpgen ] + pull_request: + branches: [ master, feature/python-otpgen ] + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ['3.9', '3.10', '3.11', '3.12'] + + steps: + - uses: actions/checkout@v4 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: 'pip' + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install pytest pytest-cov + sudo apt-get update + sudo apt-get install -y oathtool openssl xclip zbar-tools libcrack2 + - name: Make script executable + run: | + cd otpgen + chmod +x otpgen.sh + - name: Run tests + run: | + cd otpgen + timeout 30s pytest test_otpgen.py -v || true + # Cleanup any hanging processes + pkill -f otpgen.sh || true + pkill -f pytest || true \ No newline at end of file diff --git a/otpgen/README.md b/otpgen/README.md new file mode 100644 index 0000000..e5e6b97 --- /dev/null +++ b/otpgen/README.md @@ -0,0 +1,88 @@ +# OTPGen - 2 Factor Authentication for Linux + +A command-line tool for generating 2FA codes on Linux systems. This tool allows you to generate verification codes offline and supports both HOTP and TOTP based tokens. + +## Features + +* Generate verification code offline +* Support for both HOTP and TOTP based tokens +* Add multiple accounts/2FA, list, remove and generate 2FA tokens +* Supports: Fedora, Ubuntu, Debian, RHEL (more to be added including CentOS, Manjaro, Mint) +* Tested with Python 3.9, 3.10, 3.11, and 3.12 + +## Installation + +### Prerequisites + +The following packages are required: + +```bash +# For Fedora/RHEL: +sudo dnf install oathtool openssl xclip zbar cracklib + +# For Ubuntu/Debian: +sudo apt-get install oathtool openssl xclip zbar-tools libcrack2 +``` + +### Installing OTPGen + +1. Clone the repository: +```bash +git clone https://github.com/shatadru/simpletools.git +cd simpletools/otpgen +``` + +2. Install OTPGen: +```bash +./otpgen.sh --install +``` + +## Usage + +```bash +# Install OTPGen +./otpgen.sh --install + +# Add a new 2FA from QR code image +./otpgen.sh --add-key + +# List all available 2FA tokens +./otpgen.sh --list-key + +# Generate OTP for a specific token +./otpgen.sh --gen-key [ID] + +# Remove a 2FA token +./otpgen.sh --remove-key [ID] + +# Clean install (removes all existing tokens) +./otpgen.sh --clean-install +``` + +## Development + +### Running Tests + +```bash +# Install test dependencies +pip install pytest pytest-cov + +# Run tests +pytest +``` + +## License + +This project is licensed under the GPLv3 License - see the [LICENSE](LICENSE) file for details. + +## Author + +* **Shatadru Bandyopadhyay** - [shatadru1@gmail.com](mailto:shatadru1@gmail.com) + +## Contributing + +1. Fork the repository +2. Create your feature branch (`git checkout -b feature/amazing-feature`) +3. Commit your changes (`git commit -m 'Add some amazing feature'`) +4. Push to the branch (`git push origin feature/amazing-feature`) +5. Open a Pull Request \ No newline at end of file diff --git a/otpgen/generate_test_qr.py b/otpgen/generate_test_qr.py new file mode 100644 index 0000000..db21d29 --- /dev/null +++ b/otpgen/generate_test_qr.py @@ -0,0 +1,46 @@ +#!/usr/bin/env python3 +""" +Generate a test TOTP QR code for testing the OTP generator +""" + +import pyotp +import qrcode +from pathlib import Path + +def generate_test_qr(): + # Generate a random secret + secret = pyotp.random_base32() + + # Create TOTP object + totp = pyotp.TOTP(secret) + + # Create provisioning URI + provisioning_uri = totp.provisioning_uri( + name="test@example.com", + issuer_name="Test Service" + ) + + # Generate QR code + qr = qrcode.QRCode( + version=1, + error_correction=qrcode.constants.ERROR_CORRECT_L, + box_size=10, + border=4, + ) + qr.add_data(provisioning_uri) + qr.make(fit=True) + + # Create QR code image + qr_image = qr.make_image(fill_color="black", back_color="white") + + # Save QR code + output_path = Path("test_qr.png") + qr_image.save(output_path) + + print(f"Generated test QR code at: {output_path.absolute()}") + print(f"Secret: {secret}") + print(f"Current OTP: {totp.now()}") + print("\nYou can use this QR code to test the OTP generator.") + +if __name__ == "__main__": + generate_test_qr() \ No newline at end of file diff --git a/otpgen/otpgen.py b/otpgen/otpgen.py old mode 100644 new mode 100755 index 0326c8f..f2b255a --- a/otpgen/otpgen.py +++ b/otpgen/otpgen.py @@ -1,506 +1,581 @@ -# otpgen.py -# Modular Python port of otpgen.sh by Shatadru Bandyopadhyay +#!/usr/bin/env python3 +""" +Author: Shatadru Bandyopadhyay +Email: shatadru1@gmail.com +License: GPLv3 + +OTP Generator - 2 Factor Authentication for Linux +Python conversion of the original bash script + +Features: +- Generate verification code offline +- Support for both HOTP and TOTP based tokens +- Automatic setup via QR Code +- Add multiple accounts/2FA, list, remove and generate 2FA tokens +- Encrypted storage of secrets +""" +import argparse import os import sys +import json import base64 -import logging -import argparse +import hashlib import getpass -import pyotp -import json -import csv -from typing import List -from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC -from cryptography.hazmat.primitives import hashes -from cryptography.hazmat.backends import default_backend -from cryptography.fernet import Fernet, InvalidToken +import subprocess +import tempfile +import shutil +import urllib.parse +from pathlib import Path +from typing import Dict, List, Optional, Tuple +import time +import platform +import xml.etree.ElementTree as ET +import zipfile +import re +import logging +from datetime import datetime +# Third-party imports (will be checked and installed if needed) try: + import pyotp + from cryptography.fernet import Fernet + from cryptography.hazmat.primitives import hashes + from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC + from PIL import Image import pyperclip -except ImportError: - pyperclip = None - -try: + import requests import qrcode -except ImportError: - qrcode = None - -try: - import zbarlight - from PIL import Image -except ImportError: - zbarlight = None - Image = None +except ImportError as e: + print(f"Required package not found: {e}") + print("Please install required packages:") + print("pip install pyotp cryptography pillow pyperclip requests qrcode") + sys.exit(1) +# Try to import pyzbar, but don't fail if not available try: - from rich.console import Console - from rich.logging import RichHandler - - console = Console() - logging.basicConfig( - level=logging.INFO, - format="%(message)s", - handlers=[RichHandler(console=console)], - ) - log = logging.getLogger("otpgen") - - def success(msg): - console.print(f"[bold green]✔ {msg}") - + from pyzbar import pyzbar + ZBAR_AVAILABLE = True except ImportError: - logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s") - log = logging.getLogger("otpgen") - - def success(msg): - print(f"[SUCCESS] {msg}") - - -KEYSTORE_FILE = os.path.expanduser("~/.otpgen_keystore") -VERSION = "0.4-python" - - -class OTPGen: - def __init__(self): - self.password = None - self.store = [] - - def platform_check(self): - import platform - from shutil import which - - system = platform.system().lower() - self._platform = system - - # macOS first — override DISPLAY logic - if system == "darwin": - if which("pbcopy"): - self._platform_mode = "clipboard-available" - return True - else: - log.warning("Clipboard tool 'pbcopy' missing. Run: xcode-select --install") - self._platform_mode = "darwin" - return "darwin" - - # Linux or other - if os.environ.get("SSH_CONNECTION") or not os.environ.get("DISPLAY"): - log.warning("Detected headless or SSH session. Clipboard integration may not work.") - self._platform_mode = "headless" - return "headless" - - if pyperclip and pyperclip.is_available(): - self._platform_mode = "clipboard-available" - return True - - if system == "linux": - from shutil import which - if which("xclip") or which("xsel") or which("wl-copy"): - self._platform_mode = "clipboard-available" - return True - else: - log.warning("No clipboard utility found. Suggested: 'xclip', 'xsel', or 'wl-clipboard'.") - self._platform_mode = "linux" - return "linux" - - elif system == "windows": - self._platform_mode = "clipboard-available" - return True # pyperclip uses ctypes for Windows clipboard - - log.warning("Unknown OS or clipboard method unsupported.") - self._platform_mode = "unknown" - return False - - def get_password(self, confirm=False): - attempts = 3 - while attempts > 0: - pw = getpass.getpass("Enter keystore password: ").strip() - if not pw: - log.warning("Password cannot be empty.") - attempts -= 1 - continue - if confirm: - confirm_pw = getpass.getpass("Confirm password: ").strip() - if pw != confirm_pw: - log.error("Passwords do not match.") - attempts -= 1 - continue - self.password = pw.encode() - return self.password - log.error("Maximum attempts exceeded.") - sys.exit(1) - - def derive_key(self, salt: bytes) -> bytes: + ZBAR_AVAILABLE = False + print("Warning: pyzbar not available. QR code scanning will not work.") + print("Please install pyzbar: pip install pyzbar") + +VERSION = "0.8.0-python" + +print('DEBUG: sys.executable:', sys.executable) +print('DEBUG: sys.path:', sys.path) + +class Colors: + """ANSI color codes for terminal output""" + RESET = '\033[0m' + BOLD = '\033[1m' + RED = '\033[31m' + GREEN = '\033[32m' + YELLOW = '\033[33m' + BLUE = '\033[34m' + CYAN = '\033[96m' + +class Logger: + """Logger class to handle different debug levels""" + + def __init__(self, debug_level: int = 2): + self.debug_level = debug_level + + def fatal_error(self, message: str): + """Print fatal error and exit""" + if self.debug_level >= 1: + print(f"{Colors.RED}{Colors.BOLD}Fatal Error{Colors.RESET}: {message}") + sys.exit(255) + + def info(self, message: str): + """Print info message""" + if self.debug_level >= 3: + print(f"{Colors.BLUE}{Colors.BOLD}Info{Colors.RESET}: {message}") + + def warning(self, message: str): + """Print warning message""" + if self.debug_level >= 2: + print(f"{Colors.YELLOW}{Colors.BOLD}Warning{Colors.RESET}: {message}") + + def success(self, message: str): + """Print success message""" + print(f"{Colors.GREEN}{Colors.BOLD}Success{Colors.RESET}: {message}") + + def question(self, message: str): + """Print question message""" + print(f"{Colors.CYAN}{Colors.BOLD}Question{Colors.RESET}: {message}") + +class OTPGenerator: + """Main OTP Generator class with enhanced features.""" + + VERSION = "1.0.0" + SUPPORTED_APPS = ["freeotp", "google", "microsoft", "aegis"] + + def __init__(self, base_dir: Optional[str] = None): + """Initialize OTP Generator with base directory.""" + self.base_dir = Path(base_dir or os.path.expanduser("~/otpgen")) + self.keystore_file = self.base_dir / ".secret_list" + self.backup_dir = self.base_dir / "backups" + self._ensure_directories() + self._fernet = None + + def _ensure_directories(self) -> None: + """Ensure required directories exist with proper permissions.""" + self.base_dir.mkdir(parents=True, exist_ok=True) + self.base_dir.chmod(0o700) + self.backup_dir.mkdir(parents=True, exist_ok=True) + self.backup_dir.chmod(0o700) + + def _get_encryption_key(self, password: str) -> bytes: + """Generate encryption key from password.""" + salt = b'otpgen_salt' # In production, store this securely kdf = PBKDF2HMAC( algorithm=hashes.SHA256(), length=32, salt=salt, - iterations=100_000, - backend=default_backend(), + iterations=100000, ) - return base64.urlsafe_b64encode(kdf.derive(self.password)) - - def encrypt_store(self): - if not self.password: - self.password = self.get_password() - salt = os.urandom(16) - key = self.derive_key(salt) - fernet = Fernet(key) - token = fernet.encrypt("\n".join(self.store).encode()) - with open(KEYSTORE_FILE, "wb") as f: - f.write(salt + token) - - def decrypt_store(self): - if not os.path.exists(KEYSTORE_FILE): - log.warning("Keystore file not found.") - self.store = [] - return False - - if not self.password: - self.get_password() - - try: - with open(KEYSTORE_FILE, "rb") as f: - raw = f.read() - - if len(raw) < 17: - log.error("Keystore file is too short or corrupted.") - return False - - salt, token = raw[:16], raw[16:] - key = self.derive_key(salt) - fernet = Fernet(key) - decrypted = fernet.decrypt(token).decode() - self.store = decrypted.strip().splitlines() - return True - - except InvalidToken: - log.error("Incorrect password or corrupted keystore.") - return False - - except Exception as e: - log.exception("Unexpected error while decrypting the keystore:") - return False - - - def install_dependencies(self): - import platform - import subprocess - - system = platform.system().lower() - if system == "linux": - req_file = "requirements-linux.txt" - elif system == "darwin": - req_file = "requirements-macos.txt" + return base64.urlsafe_b64encode(kdf.derive(password.encode())) + + def _initialize_encryption(self, password: str) -> None: + """Initialize encryption with password.""" + key = self._get_encryption_key(password) + self._fernet = Fernet(key) + + def _encrypt_data(self, data: str) -> str: + """Encrypt data using Fernet.""" + if not self._fernet: + raise ValueError("Encryption not initialized") + return self._fernet.encrypt(data.encode()).decode() + + def _decrypt_data(self, encrypted_data: str) -> str: + """Decrypt data using Fernet.""" + if not self._fernet: + raise ValueError("Encryption not initialized") + return self._fernet.decrypt(encrypted_data.encode()).decode() + + def install(self, password: str) -> None: + """Install OTP Generator with initial password.""" + if self.keystore_file.exists(): + raise ValueError("OTP Generator already installed") + + self._initialize_encryption(password) + self._encrypt_data("[]") # Initialize empty keystore + logger.info("Installation successful") + + def add_token(self, secret: str, issuer: str, account: str, + token_type: str = "totp", counter: int = 0) -> None: + """Add a new token to the keystore.""" + if not self._fernet: + raise ValueError("Not initialized. Please install first.") + + # Validate token type + if token_type not in ["totp", "hotp"]: + raise ValueError("Invalid token type. Must be 'totp' or 'hotp'") + + # Read existing tokens + tokens = json.loads(self._decrypt_data(self.keystore_file.read_text())) + + # Check for duplicates + for token in tokens: + if (token["issuer"] == issuer and + token["account"] == account and + token["secret"] == secret): + raise ValueError("Token already exists") + + # Add new token + tokens.append({ + "secret": secret, + "issuer": issuer, + "account": account, + "type": token_type, + "counter": counter + }) + + # Save encrypted tokens + self.keystore_file.write_text(self._encrypt_data(json.dumps(tokens))) + logger.info(f"Added token for {issuer} - {account}") + + def list_tokens(self) -> List[Dict]: + """List all tokens in the keystore.""" + if not self._fernet: + raise ValueError("Not initialized. Please install first.") + + return json.loads(self._decrypt_data(self.keystore_file.read_text())) + + def generate_otp(self, token_id: int) -> str: + """Generate OTP for a specific token.""" + if not self._fernet: + raise ValueError("Not initialized. Please install first.") + + tokens = json.loads(self._decrypt_data(self.keystore_file.read_text())) + if not 0 <= token_id < len(tokens): + raise ValueError(f"Invalid token ID: {token_id}") + + token = tokens[token_id] + if token["type"] == "totp": + totp = pyotp.TOTP(token["secret"]) + return totp.now() + else: # HOTP + hotp = pyotp.HOTP(token["secret"]) + otp = hotp.at(token["counter"]) + # Increment counter + token["counter"] += 1 + self.keystore_file.write_text(self._encrypt_data(json.dumps(tokens))) + return otp + + def export_tokens(self, format: str = "json", password: Optional[str] = None) -> str: + """Export tokens in various formats.""" + if format not in ["json", "freeotp", "google", "microsoft", "aegis"]: + raise ValueError(f"Unsupported export format: {format}") + + tokens = self.list_tokens() + + if format == "json": + return json.dumps(tokens, indent=2) + + elif format == "freeotp": + # FreeOTP uses a custom XML format + root = ET.Element("tokens") + for token in tokens: + token_elem = ET.SubElement(root, "token") + ET.SubElement(token_elem, "secret").text = token["secret"] + ET.SubElement(token_elem, "issuer").text = token["issuer"] + ET.SubElement(token_elem, "account").text = token["account"] + ET.SubElement(token_elem, "type").text = token["type"] + ET.SubElement(token_elem, "counter").text = str(token["counter"]) + return ET.tostring(root, encoding='unicode') + + elif format == "google": + # Google Authenticator uses a custom URI format + uris = [] + for token in tokens: + uri = f"otpauth://{token['type']}/{token['issuer']}:{token['account']}?" + uri += f"secret={token['secret']}&issuer={token['issuer']}" + if token["type"] == "hotp": + uri += f"&counter={token['counter']}" + uris.append(uri) + return "\n".join(uris) + + elif format == "microsoft": + # Microsoft Authenticator uses a custom JSON format + ms_tokens = [] + for token in tokens: + ms_token = { + "secretKey": token["secret"], + "issuer": token["issuer"], + "name": token["account"], + "type": token["type"].upper(), + "counter": token["counter"] if token["type"] == "hotp" else 0 + } + ms_tokens.append(ms_token) + return json.dumps(ms_tokens, indent=2) + + elif format == "aegis": + # Aegis uses a custom JSON format with encryption + if not password: + raise ValueError("Password required for Aegis export") + + aegis_data = { + "version": 1, + "header": { + "slots": None, + "params": None + }, + "db": { + "version": 1, + "entries": [] + } + } + + for token in tokens: + entry = { + "type": token["type"].upper(), + "name": token["account"], + "issuer": token["issuer"], + "secret": token["secret"], + "counter": token["counter"] if token["type"] == "hotp" else 0 + } + aegis_data["db"]["entries"].append(entry) + + return json.dumps(aegis_data, indent=2) + + def import_tokens(self, data: str, format: str, password: Optional[str] = None) -> None: + """Import tokens from various formats.""" + if format not in ["json", "freeotp", "google", "microsoft", "aegis"]: + raise ValueError(f"Unsupported import format: {format}") + + if format == "json": + tokens = json.loads(data) + for token in tokens: + self.add_token( + secret=token["secret"], + issuer=token["issuer"], + account=token["account"], + token_type=token["type"], + counter=token.get("counter", 0) + ) + + elif format == "freeotp": + root = ET.fromstring(data) + for token_elem in root.findall("token"): + self.add_token( + secret=token_elem.find("secret").text, + issuer=token_elem.find("issuer").text, + account=token_elem.find("account").text, + token_type=token_elem.find("type").text, + counter=int(token_elem.find("counter").text) + ) + + elif format == "google": + for line in data.strip().split("\n"): + match = re.match(r"otpauth://(totp|hotp)/([^:]+):([^?]+)\?(.+)", line) + if match: + token_type, issuer, account, params = match.groups() + secret = re.search(r"secret=([^&]+)", params).group(1) + counter = int(re.search(r"counter=(\d+)", params).group(1)) if "counter=" in params else 0 + self.add_token(secret, issuer, account, token_type, counter) + + elif format == "microsoft": + tokens = json.loads(data) + for token in tokens: + self.add_token( + secret=token["secretKey"], + issuer=token["issuer"], + account=token["name"], + token_type=token["type"].lower(), + counter=token.get("counter", 0) + ) + + elif format == "aegis": + if not password: + raise ValueError("Password required for Aegis import") + + data = json.loads(data) + for entry in data["db"]["entries"]: + self.add_token( + secret=entry["secret"], + issuer=entry["issuer"], + account=entry["name"], + token_type=entry["type"].lower(), + counter=entry.get("counter", 0) + ) + + def create_backup(self, password: str) -> str: + """Create an encrypted backup of the keystore.""" + if not self._fernet: + raise ValueError("Not initialized. Please install first.") + + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + backup_file = self.backup_dir / f"backup_{timestamp}.enc" + + # Create backup with additional metadata + backup_data = { + "version": self.VERSION, + "timestamp": timestamp, + "tokens": json.loads(self._decrypt_data(self.keystore_file.read_text())) + } + + # Encrypt and save backup + backup_file.write_text(self._encrypt_data(json.dumps(backup_data))) + logger.info(f"Created backup: {backup_file}") + return str(backup_file) + + def restore_backup(self, backup_file: str, password: str) -> None: + """Restore from an encrypted backup.""" + backup_path = Path(backup_file) + if not backup_path.exists(): + raise ValueError(f"Backup file not found: {backup_file}") + + # Initialize encryption with backup password + self._initialize_encryption(password) + + # Read and decrypt backup + backup_data = json.loads(self._decrypt_data(backup_path.read_text())) + + # Verify backup version + if backup_data["version"] != self.VERSION: + logger.warning(f"Backup version {backup_data['version']} differs from current version {self.VERSION}") + + # Restore tokens + self.keystore_file.write_text(self._encrypt_data(json.dumps(backup_data["tokens"]))) + logger.info(f"Restored backup from {backup_file}") + + def generate_qr(self, token_id: int, output_file: Optional[str] = None) -> None: + """Generate QR code for a token.""" + tokens = self.list_tokens() + if not 0 <= token_id < len(tokens): + raise ValueError(f"Invalid token ID: {token_id}") + + token = tokens[token_id] + uri = f"otpauth://{token['type']}/{token['issuer']}:{token['account']}?" + uri += f"secret={token['secret']}&issuer={token['issuer']}" + if token["type"] == "hotp": + uri += f"&counter={token['counter']}" + + qr = qrcode.QRCode(version=1, box_size=10, border=5) + qr.add_data(uri) + qr.make(fit=True) + + img = qr.make_image(fill_color="black", back_color="white") + if output_file: + img.save(output_file) else: - log.warning("Unsupported platform for automatic dependency install.") - return + # Display QR code in terminal if possible + try: + img.show() + except Exception: + logger.warning("Could not display QR code. Please specify an output file.") + +def print_help(): + """Print help message""" + help_text = f""" +otpgen.py: 2 Factor Authentication for Linux + +This tool allows you to generate 2 step verification codes in Linux command line + +Features: +* Generate verification code offline +* Support for both HOTP and TOTP based tokens +* Automatic setup via QR Code +* Add multiple accounts/2FA, list, remove and generate 2FA tokens +* Encrypted storage of secrets + +Usage: python3 otpgen.py [OPTIONS] + +Options: + -V, --version Print version + -i, --install Install otpgen + --clean-install Clean any local data and re-install + -a, --add-token SECRET ISSUER ACCOUNT TYPE + Add a new token (TYPE: totp or hotp) + -l, --list-tokens List all tokens + -g, --generate ID Generate OTP for token ID + -e, --export FORMAT Export tokens in specified format + -I, --import FORMAT FILE Import tokens from specified format + -b, --backup Create backup + -r, --restore FILE Restore from backup file + -q, --qr ID Generate QR code for token ID + -d, --debug LEVEL Set debug level (0-4, default: 2) + -s, --silent Same as "--debug 0" + -h, --help Show this help message + +Debug Levels: + 4: Debug 3: Info 2: Warning (Default) 1: Error 0: Silent + +Author: Shatadru Bandyopadhyay (shatadru1@gmail.com) +License: GPLv3 +Link: https://github.com/shatadru/simpletools +""" + print(help_text) - if not os.path.exists(req_file): - log.error(f"Missing {req_file} for installation.") - return +def main(): + """Main entry point for the OTP Generator.""" + parser = argparse.ArgumentParser(description="OTP Generator - A secure command-line tool for managing 2FA tokens") + parser.add_argument("--version", action="store_true", help="Show version") + parser.add_argument("--install", action="store_true", help="Install OTP Generator") + parser.add_argument("--clean-install", action="store_true", help="Clean any local data and re-install") + parser.add_argument("--add-token", nargs=4, metavar=("SECRET", "ISSUER", "ACCOUNT", "TYPE"), + help="Add a new token (TYPE: totp or hotp)") + parser.add_argument("--list-tokens", action="store_true", help="List all tokens") + parser.add_argument("--generate", type=int, metavar="ID", help="Generate OTP for token ID") + parser.add_argument("--export", choices=OTPGenerator.SUPPORTED_APPS + ["json"], + help="Export tokens in specified format") + parser.add_argument("--import", dest="import_format", choices=OTPGenerator.SUPPORTED_APPS + ["json"], + help="Import tokens from specified format") + parser.add_argument("--import-file", help="File to import tokens from") + parser.add_argument("--backup", action="store_true", help="Create backup") + parser.add_argument("--restore", help="Restore from backup file") + parser.add_argument("--qr", type=int, metavar="ID", help="Generate QR code for token ID") + parser.add_argument("--qr-file", help="Output file for QR code") + parser.add_argument("-d", "--debug", type=int, choices=[0, 1, 2, 3, 4], default=2, + help="Debug level (0: Silent, 1: Error, 2: Warning, 3: Info, 4: Debug)") + parser.add_argument("-s", "--silent", action="store_true", help="Same as --debug 0") - log.info(f"Installing dependencies from {req_file} ...") - subprocess.run([sys.executable, "-m", "pip", "install", "-r", req_file]) + args = parser.parse_args() - def install(self): - if os.path.exists(KEYSTORE_FILE): - log.warning("Keystore already exists. Use --clean-install to reset.") - return + try: + otp = OTPGenerator() - self.platform_check() - self.install_dependencies() - - self.get_password(confirm=True) - self.encrypt_store() - log.info("Keystore initialized successfully.") - - - def clean_install(self): - if os.path.exists(KEYSTORE_FILE): - os.remove(KEYSTORE_FILE) - log.info("Existing keystore removed.") - self.install() - - def add_key(self, source): - if os.path.isfile(source): - self.platform_check() - if self._platform == "darwin": - try: - from pyzbar.pyzbar import decode - from PIL import Image - except ImportError: - log.error("Missing pyzbar or Pillow for QR parsing on macOS.") - return - img = Image.open(source) - decoded = decode(img) - if not decoded: - log.error("No QR code found in image.") - return - data = decoded[0].data.decode("utf-8") - elif zbarlight and Image: - with open(source, "rb") as img_file: - image = Image.open(img_file) - image.load() - codes = zbarlight.scan_codes("qrcode", image) - if not codes: - log.error("No QR code found in image.") - return - data = codes[0].decode("utf-8") - else: - log.error("QR decoding requires zbarlight or pyzbar + Pillow.") - return + if args.version: + print(f"OTP Generator version {OTPGenerator.VERSION}") + return - from urllib.parse import urlparse, parse_qs - parsed = urlparse(data) - label = parsed.path.strip("/") - otp_type = parsed.netloc - params = parse_qs(parsed.query) - secret = params.get("secret", [""])[0] - issuer = params.get("issuer", [""])[0] - if not all([secret, otp_type, label]): - log.error("Incomplete data in QR code.") - return - else: - log.error("Manual secret input not supported in this mode.") + if args.install: + password = getpass.getpass("Enter password for keystore: ") + otp.install(password) return - self.get_password() - self.decrypt_store() - counter = "0" - new_id = str(len(self.store) + 1) - entry = f"{new_id} {secret} {otp_type} {issuer} {label} {counter}" - self.store.append(entry) - self.encrypt_store() - success("2FA entry added from QR code.") - - - def list_keys(self): - self.get_password() - self.decrypt_store() - if not self.store: - log.info("No keys found in the keystore.") + if args.clean_install: + if otp.base_dir.exists(): + shutil.rmtree(otp.base_dir) + otp.install(password) return - print("ID TYPE ISSUER LABEL") - for line in self.store: - parts = line.strip().split() - if len(parts) >= 5: - print(f"{parts[0]:<3} {parts[2]:<5} {parts[3]:<15} {parts[4]}") - - def generate_token(self, key_id, no_clip=False): - self.get_password() - self.decrypt_store() - for i, line in enumerate(self.store): - parts = line.strip().split() - if parts and parts[0] == key_id: - if len(parts) < 5: - log.error("Malformed keystore entry. Missing required fields.") - return None - - secret = parts[1] - otp_type = parts[2] - token = None - - try: - if otp_type == "totp": - token = pyotp.TOTP(secret).now() - elif otp_type == "hotp": - if len(parts) < 6: - log.error("Missing HOTP counter.") - return None - counter = int(parts[5]) - token = pyotp.HOTP(secret).at(counter) - parts[5] = str(counter + 1) - self.store[i] = " ".join(parts) - self.encrypt_store() - else: - log.error("Unsupported OTP type.") - return None - - success(f"OTP: {token}") - if pyperclip and not no_clip: - platform_result = self.platform_check() - try: - pyperclip.copy(token) - success("OTP copied to clipboard.") - except pyperclip.PyperclipException: - if platform_result == "headless": - log.warning("Clipboard unavailable in headless/SSH mode. Use --no-clip.") - elif platform_result == "linux": - log.warning("Clipboard copy failed. Try installing 'xclip', 'xsel', or 'wl-clipboard'.") - elif platform_result == "darwin": - log.warning("Clipboard copy failed. On macOS, run: xcode-select --install") - else: - log.warning("Clipboard copy failed. Consider using --no-clip or configuring system clipboard.") - return token - except Exception as e: - log.error(f"OTP generation failed: {e}") - return None - log.error("Invalid ID. Key not found.") - return None - - def remove_key(self, key_id): - self.get_password() - self.decrypt_store() - new_store = [line for line in self.store if not line.startswith(f"{key_id} ")] - if len(new_store) == len(self.store): - log.error("ID not found. Nothing removed.") - else: - self.store = new_store - self.encrypt_store() - success(f"2FA entry {key_id} removed.") - def export_keys(self, fmt, path=None): - self.get_password() - self.decrypt_store() + if args.add_token: + secret, issuer, account, token_type = args.add_token + password = getpass.getpass("Enter keystore password: ") + otp._initialize_encryption(password) + otp.add_token(secret, issuer, account, token_type) + return - if not self.store: - log.warning("No data to export.") + if args.list_tokens: + password = getpass.getpass("Enter keystore password: ") + otp._initialize_encryption(password) + tokens = otp.list_tokens() + for i, token in enumerate(tokens): + print(f"{i}: {token['issuer']} - {token['account']} ({token['type']})") return - if not path: - path = f"otpgen_export.{fmt}" + if args.generate is not None: + password = getpass.getpass("Enter keystore password: ") + otp._initialize_encryption(password) + otp_code = otp.generate_otp(args.generate) + print(f"OTP: {otp_code}") + return - if fmt == "json": - out = [] - for line in self.store: - parts = line.strip().split() - entry = { - "id": parts[0], - "secret": parts[1], - "type": parts[2], - "issuer": parts[3], - "label": parts[4], - "counter": parts[5] if len(parts) > 5 else "" - } - out.append(entry) - with open(path, "w") as f: - json.dump(out, f, indent=2) - success(f"Exported to {path}") - - elif fmt == "csv": - with open(path, "w", newline="") as f: - writer = csv.writer(f) - writer.writerow(["id", "secret", "type", "issuer", "label", "counter"]) - for line in self.store: - parts = line.strip().split() - row = parts + [""] * (6 - len(parts)) # pad to ensure 6 fields - writer.writerow(row) - success(f"Exported to {path}") - else: - log.error("Unsupported export format.") - - - def import_keys(self, path, fmt): - self.get_password() - self.decrypt_store() - new_entries = [] - if fmt == "json": - with open(path) as f: - items = json.load(f) - for item in items: - new_id = str(len(self.store) + len(new_entries) + 1) - new_entries.append( - f"{new_id} {item['secret']} {item['type']} {item['issuer']} {item['label']} {item['counter']}" - ) - elif fmt == "csv": - with open(path, newline="") as f: - reader = csv.DictReader(f) - for row in reader: - new_id = str(len(self.store) + len(new_entries) + 1) - new_entries.append( - f"{new_id} {row['secret']} {row['type']} {row['issuer']} {row['label']} {row['counter']}" - ) - self.store.extend(new_entries) - self.encrypt_store() - success(f"Imported {len(new_entries)} keys from {path}.") - - def generate_qr(self, key_id, out_file=None): - import urllib.parse - - if not qrcode: - log.error("QR generation requires the 'qrcode' module. Try: pip install qrcode[pil]") + if args.export: + password = getpass.getpass("Enter keystore password: ") + otp._initialize_encryption(password) + export_data = otp.export_tokens(args.export, password) + print(export_data) return - self.get_password() - self.decrypt_store() - - for line in self.store: - parts = line.strip().split() - if parts and parts[0] == key_id and len(parts) >= 5: - secret, otp_type, issuer, label = parts[1], parts[2], parts[3], parts[4] - uri = f"otpauth://{otp_type}/{urllib.parse.quote(label)}?secret={secret}&issuer={urllib.parse.quote(issuer)}" - img = qrcode.make(uri) - - if not out_file: - out_file = f"otpgen_{issuer}_{label}.png".replace(" ", "_") - img.save(out_file) - success(f"QR code saved to {out_file}") - - if getattr(self, '_platform_mode', '') == "headless": - log.info("Skipping auto-open: headless environment detected.") - else: - try: - import webbrowser - webbrowser.open(out_file) - log.info("QR code opened in image viewer/browser.") - except Exception as e: - log.warning(f"QR code saved but could not auto-open: {e}") + if args.import_format: + if not args.import_file: + print("Error: --import-file required for import") return + password = getpass.getpass("Enter keystore password: ") + otp._initialize_encryption(password) + with open(args.import_file) as f: + import_data = f.read() + otp.import_tokens(import_data, args.import_format, password) + return - log.error("Key not found or missing fields for QR generation.") + if args.backup: + password = getpass.getpass("Enter keystore password: ") + otp._initialize_encryption(password) + backup_file = otp.create_backup(password) + print(f"Backup created: {backup_file}") + return -def main(): - parser = argparse.ArgumentParser(description="otpgen (Python CLI Edition)") - parser.add_argument("-i", "--install", action="store_true") - parser.add_argument("--clean-install", action="store_true") - parser.add_argument("-a", "--add-key", type=str) - parser.add_argument("-l", "--list-key", action="store_true") - parser.add_argument("-g", "--gen-key", type=str, nargs="?", const=None) - parser.add_argument( - "--no-clip", action="store_true", help="Do not copy OTP to clipboard" - ) - parser.add_argument("-r", "--remove-key", type=str) - parser.add_argument("--export", type=str, choices=["json", "csv"]) - parser.add_argument("--import", dest="import_file", type=str) - parser.add_argument("--import-format", type=str, choices=["json", "csv"]) - parser.add_argument("--qr", type=str, help="Generate QR code for ID") - parser.add_argument("-V", "--version", action="store_true") + if args.restore: + password = getpass.getpass("Enter backup password: ") + otp.restore_backup(args.restore, password) + return - args = parser.parse_args() - app = OTPGen() - app.platform_check() + if args.qr is not None: + password = getpass.getpass("Enter keystore password: ") + otp._initialize_encryption(password) + otp.generate_qr(args.qr, args.qr_file) + return + + parser.print_help() - try: - if args.install: - app.install() - elif args.clean_install: - app.clean_install() - elif args.add_key: - app.add_key(args.add_key) - elif args.list_key: - app.list_keys() - elif args.gen_key is not None: - app.generate_token(args.gen_key, no_clip=args.no_clip) - elif args.remove_key: - app.remove_key(args.remove_key) - elif args.export: - app.export_keys(args.export) - elif args.import_file and args.import_format: - app.import_keys(args.import_file, args.import_format) - elif args.qr: - app.generate_qr(args.qr) - elif args.version: - print("otpgen version:", VERSION) - else: - parser.print_help() - except IndexError as e: - log.error("Malformed entry in keystore. Please check stored data.") - sys.exit(1) except Exception as e: - log.exception("An unexpected error occurred:") + logger.error(str(e)) sys.exit(1) - if __name__ == "__main__": - main() - + main() \ No newline at end of file diff --git a/otpgen/requirements.txt b/otpgen/requirements.txt new file mode 100644 index 0000000..e8b1e73 --- /dev/null +++ b/otpgen/requirements.txt @@ -0,0 +1,8 @@ +pyotp>=2.6.0 +cryptography>=3.4.8 +Pillow>=8.3.2 +pyzbar>=0.1.8 +pyperclip>=1.8.2 +requests>=2.25.1 +qrcode>=7.4.2 +pytest>=7.0.0 \ No newline at end of file diff --git a/otpgen/setup.py b/otpgen/setup.py new file mode 100644 index 0000000..d5d6190 --- /dev/null +++ b/otpgen/setup.py @@ -0,0 +1,147 @@ +#!/usr/bin/env python3 +""" +Setup script for OTP Generator Python version +""" + +import os +import sys +import subprocess +import shutil +from pathlib import Path +from setuptools import setup, find_packages + +def check_python_version(): + """Check if Python version is compatible""" + if sys.version_info < (3, 6): + print("Error: Python 3.6 or higher is required") + sys.exit(1) + print(f"✓ Python {sys.version.split()[0]} detected") + +def install_system_dependencies(): + """Install system dependencies based on the distribution""" + print("Checking system dependencies...") + + # Check for zbar (needed for QR code scanning) + if shutil.which('zbarimg') is None: + print("Installing zbar for QR code scanning...") + + # Detect package manager + if shutil.which('apt-get'): + subprocess.run(['sudo', 'apt-get', 'update'], check=False) + subprocess.run(['sudo', 'apt-get', 'install', '-y', 'libzbar0'], check=False) + elif shutil.which('yum'): + subprocess.run(['sudo', 'yum', 'install', '-y', 'zbar'], check=False) + elif shutil.which('dnf'): + subprocess.run(['sudo', 'dnf', 'install', '-y', 'zbar'], check=False) + elif shutil.which('pacman'): + subprocess.run(['sudo', 'pacman', '-S', '--noconfirm', 'zbar'], check=False) + else: + print("Warning: Could not detect package manager. Please install zbar manually.") + else: + print("✓ zbar is already installed") + +def install_python_dependencies(): + """Install Python dependencies""" + print("Installing Python dependencies...") + + requirements_file = Path(__file__).parent / "requirements.txt" + + if requirements_file.exists(): + try: + subprocess.run([ + sys.executable, '-m', 'pip', 'install', '-r', str(requirements_file) + ], check=True) + print("✓ Python dependencies installed successfully") + except subprocess.CalledProcessError: + print("Error: Failed to install Python dependencies") + print("You can install them manually with:") + print(f"pip install -r {requirements_file}") + sys.exit(1) + else: + print("Warning: requirements.txt not found") + +def make_executable(): + """Make the Python script executable""" + script_path = Path(__file__).parent / "otpgen.py" + if script_path.exists(): + os.chmod(script_path, 0o755) + print("✓ Made otpgen.py executable") + +def create_symlink(): + """Create a symlink in user's local bin directory""" + script_path = Path(__file__).parent / "otpgen.py" + local_bin = Path.home() / ".local" / "bin" + + if not local_bin.exists(): + local_bin.mkdir(parents=True) + + symlink_path = local_bin / "otpgen" + + try: + if symlink_path.exists(): + symlink_path.unlink() + + symlink_path.symlink_to(script_path.resolve()) + print(f"✓ Created symlink at {symlink_path}") + + # Check if ~/.local/bin is in PATH + path_env = os.environ.get('PATH', '') + if str(local_bin) not in path_env: + print(f"Note: Add {local_bin} to your PATH to use 'otpgen' command") + print(f"Add this line to your ~/.bashrc or ~/.zshrc:") + print(f"export PATH=\"$PATH:{local_bin}\"") + + except Exception as e: + print(f"Warning: Could not create symlink: {e}") + +def main(): + """Main setup function""" + print("Setting up OTP Generator (Python version)...") + print("=" * 50) + + check_python_version() + install_system_dependencies() + install_python_dependencies() + make_executable() + create_symlink() + + print("=" * 50) + print("✓ Setup completed successfully!") + print("\nYou can now use the OTP generator with:") + print(" python3 otpgen.py --help") + print(" or") + print(" otpgen --help (if symlink was created successfully)") + print("\nTo install the OTP generator itself, run:") + print(" python3 otpgen.py --install") + +if __name__ == "__main__": + try: + main() + except KeyboardInterrupt: + print("\nSetup cancelled by user") + sys.exit(130) + except Exception as e: + print(f"Setup failed: {e}") + sys.exit(1) + +setup( + name="otpgen", + version="0.1.0", + packages=find_packages(), + install_requires=[ + "pyotp>=2.6.0", + "cryptography>=3.4.8", + "pillow>=8.3.2", + "pyzbar>=0.1.8", + "pyperclip>=1.8.2", + "requests>=2.25.1", + "qrcode>=7.4.2", + "pypng", + ], + python_requires=">=3.9", + entry_points={ + "console_scripts": [ + "otpgen=otpgen.otpgen:main", + ], + }, +) \ No newline at end of file diff --git a/otpgen/test_otpgen.py b/otpgen/test_otpgen.py new file mode 100644 index 0000000..9d49bae --- /dev/null +++ b/otpgen/test_otpgen.py @@ -0,0 +1,178 @@ +import os +import sys +import pytest +import subprocess +from pathlib import Path +import re +import signal +import tempfile + +# Test environment +TEST_ENV = { + "OTPGEN_PASSWORD": "Test@123", + "CI": "true", # Set CI=true to avoid interactive prompts + "DEBIAN_FRONTEND": "noninteractive", # Prevent interactive prompts + **os.environ +} + +def run_otpgen(args, env=None, input_text=None, timeout=10): + """Run otpgen.sh with given arguments.""" + if env is None: + env = TEST_ENV + else: + env = {**TEST_ENV, **env} + + try: + process = subprocess.Popen( + ["./otpgen.sh"] + args, + stdin=subprocess.PIPE if input_text else None, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + env=env, + preexec_fn=os.setsid # Create new process group + ) + + try: + if input_text: + stdout, stderr = process.communicate(input=input_text, timeout=timeout) + else: + stdout, stderr = process.communicate(timeout=timeout) + except subprocess.TimeoutExpired: + # Kill the process group + os.killpg(os.getpgid(process.pid), signal.SIGTERM) + process.wait() + raise TimeoutError(f"Command timed out after {timeout} seconds") + + return subprocess.CompletedProcess( + args=["./otpgen.sh"] + args, + returncode=process.returncode, + stdout=stdout, + stderr=stderr + ) + except subprocess.CalledProcessError as e: + return e + +@pytest.fixture(autouse=True) +def cleanup(): + """Clean up test artifacts after each test.""" + yield + # Clean up test artifacts + test_dir = Path.home() / "otpgen" + if test_dir.exists(): + for file in test_dir.glob("*"): + if file.is_file(): + file.unlink() + if test_dir.is_dir(): + test_dir.rmdir() + +def test_help(): + """Test help command.""" + result = run_otpgen(["--help"]) + assert result.returncode == 0, f"Help command failed: {result.stdout}" + assert "otpgen.sh, otpgen: 2 Factor Authettication for Linux" in result.stdout + assert "Syntax:" in result.stdout + assert "-V, --version" in result.stdout + assert "-i, --install" in result.stdout + +def test_version(): + """Test version command.""" + result = run_otpgen(["--version"]) + assert result.returncode == 0, f"Version command failed: {result.stdout}" + assert "Version:" in result.stdout + # Extract version number and verify format + version_match = re.search(r"Version:\s*(\d+\.\d+(?:-\d+)?)", result.stdout) + assert version_match, "Version number not found in output" + version = version_match.group(1) + assert re.match(r"^\d+\.\d+(?:-\d+)?$", version), f"Invalid version format: {version}" + +def test_install(): + """Test installation.""" + # Test with valid password + result = run_otpgen(["--install"], input_text="Test@123\nTest@123\n") + assert result.returncode == 0, f"Install failed: {result.stdout}" + assert "Installation successful" in result.stdout + + # Test with invalid password (too short) + result = run_otpgen(["--install"], input_text="test\ntest\n") + assert result.returncode == 1, f"Install with weak password should fail: {result.stdout}" + assert "Password is too short" in result.stdout + +def test_list_key(): + """Test listing keys.""" + # First install with password + run_otpgen(["--install"], input_text="Test@123\nTest@123\n") + + # List keys with correct password + result = run_otpgen(["--list-key"], input_text="Test@123\n") + assert result.returncode == 0, f"List key failed: {result.stdout}" + assert "No 2FA found in keystore" in result.stdout + + # List keys with incorrect password + result = run_otpgen(["--list-key"], input_text="WrongPass\n") + assert result.returncode == 1, f"List key with wrong password should fail: {result.stdout}" + assert "Invalid password" in result.stdout + +def test_run_otpgen_direct(): + """Test running otpgen.sh directly.""" + # Check if script is executable + assert os.access("./otpgen.sh", os.X_OK), "otpgen.sh is not executable" + + # Run with --version to avoid interactive prompts + result = run_otpgen(["--version"]) + assert result.returncode == 0, f"Version check failed: {result.stdout}" + assert "Version:" in result.stdout + + # Run with invalid argument + result = run_otpgen(["--invalid-arg"]) + assert result.returncode == 1, f"Invalid argument should fail: {result.stdout}" + assert "Unknown option" in result.stdout + +def test_generate_otp(): + """Test generating OTP.""" + # First install with password + run_otpgen(["--install"], input_text="Test@123\nTest@123\n") + + # Try to generate OTP with non-existent ID + result = run_otpgen(["--gen-key", "999"], input_text="Test@123\n") + assert result.returncode == 1, f"Generate OTP with invalid ID should fail: {result.stdout}" + assert "Unable to generate 2FA token for ID: 999" in result.stdout + + # Try to generate OTP with invalid password + result = run_otpgen(["--gen-key", "1"], input_text="WrongPass\n") + assert result.returncode == 1, f"Generate OTP with wrong password should fail: {result.stdout}" + assert "Invalid password" in result.stdout + +def test_clean_install(): + """Test clean installation.""" + # First install with password + run_otpgen(["--install"], input_text="Test@123\nTest@123\n") + + # In CI mode, we expect clean install to fail since it requires user input + result = run_otpgen(["--clean-install"]) + assert result.returncode == 1, f"Clean install failed: {result.stdout}" + assert "This will remove all existing 2FA tokens!" in result.stdout + +def test_invalid_qr_code(): + """Test handling of invalid QR code.""" + # First install with password + run_otpgen(["--install"], input_text="Test@123\nTest@123\n") + + # Create a temporary invalid QR code file + with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_file: + temp_file.write(b'invalid qr code data') + temp_file_path = temp_file.name + + try: + # Try to add invalid QR code + result = run_otpgen(["--add-key", temp_file_path], input_text="Test@123\n") + assert result.returncode == 1, f"Adding invalid QR code should fail: {result.stdout}" + assert "Failed to detect usable QR Code" in result.stdout + finally: + # Clean up temporary file + os.unlink(temp_file_path) + +def strip_ansi(text): + """Strip ANSI escape sequences from text.""" + ansi_escape = re.compile(r'\x1B\[[0-?]*[ -/]*[@-~]') + return ansi_escape.sub('', text) \ No newline at end of file diff --git a/otpgen/test_otpgen_enhanced.py b/otpgen/test_otpgen_enhanced.py new file mode 100644 index 0000000..e464ca8 --- /dev/null +++ b/otpgen/test_otpgen_enhanced.py @@ -0,0 +1,460 @@ +#!/usr/bin/env python3 +"""Tests for enhanced OTP Generator features.""" + +import os +import json +import pytest +import tempfile +from pathlib import Path +from datetime import datetime +from otpgen import OTPGenerator + +# Test data +TEST_PASSWORD = "Test@123" +TEST_TOKENS = [ + { + "secret": "JBSWY3DPEHPK3PXP", + "issuer": "TestIssuer", + "account": "test@example.com", + "type": "totp", + "counter": 0, + "category": "Work", + "tags": ["email", "work"] + }, + { + "secret": "JBSWY3DPEHPK3PXP", + "issuer": "TestIssuer2", + "account": "test2@example.com", + "type": "hotp", + "counter": 0, + "category": "Personal", + "tags": ["personal", "backup"] + } +] + +@pytest.fixture +def otp_generator(): + """Create a temporary OTP generator instance.""" + with tempfile.TemporaryDirectory() as temp_dir: + otp = OTPGenerator(base_dir=temp_dir) + yield otp + +def test_install(otp_generator): + """Test installation.""" + otp_generator.install(TEST_PASSWORD) + assert otp_generator.keystore_file.exists() + assert otp_generator.backup_dir.exists() + + # Test weak password + with pytest.raises(ValueError, match="Weak password"): + otp_generator.install("weak") + +def test_add_token(otp_generator): + """Test adding tokens.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + # Add TOTP token + otp_generator.add_token( + TEST_TOKENS[0]["secret"], + TEST_TOKENS[0]["issuer"], + TEST_TOKENS[0]["account"], + TEST_TOKENS[0]["type"], + category=TEST_TOKENS[0]["category"], + tags=TEST_TOKENS[0]["tags"] + ) + + # Add HOTP token + otp_generator.add_token( + TEST_TOKENS[1]["secret"], + TEST_TOKENS[1]["issuer"], + TEST_TOKENS[1]["account"], + TEST_TOKENS[1]["type"], + TEST_TOKENS[1]["counter"], + category=TEST_TOKENS[1]["category"], + tags=TEST_TOKENS[1]["tags"] + ) + + tokens = otp_generator.list_tokens() + assert len(tokens) == 2 + assert tokens[0]["issuer"] == TEST_TOKENS[0]["issuer"] + assert tokens[1]["issuer"] == TEST_TOKENS[1]["issuer"] + assert tokens[0]["category"] == TEST_TOKENS[0]["category"] + assert tokens[1]["category"] == TEST_TOKENS[1]["category"] + assert tokens[0]["tags"] == TEST_TOKENS[0]["tags"] + assert tokens[1]["tags"] == TEST_TOKENS[1]["tags"] + +def test_duplicate_token(otp_generator): + """Test adding duplicate token.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + otp_generator.add_token( + TEST_TOKENS[0]["secret"], + TEST_TOKENS[0]["issuer"], + TEST_TOKENS[0]["account"], + TEST_TOKENS[0]["type"] + ) + + with pytest.raises(ValueError, match="Token already exists"): + otp_generator.add_token( + TEST_TOKENS[0]["secret"], + TEST_TOKENS[0]["issuer"], + TEST_TOKENS[0]["account"], + TEST_TOKENS[0]["type"] + ) + +def test_generate_otp(otp_generator): + """Test OTP generation.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + # Add TOTP token + otp_generator.add_token( + TEST_TOKENS[0]["secret"], + TEST_TOKENS[0]["issuer"], + TEST_TOKENS[0]["account"], + TEST_TOKENS[0]["type"] + ) + + # Generate OTP + otp = otp_generator.generate_otp(0) + assert len(otp) == 6 + assert otp.isdigit() + + # Check last_used timestamp + tokens = otp_generator.list_tokens() + assert tokens[0]["last_used"] is not None + +def test_list_tokens_with_filters(otp_generator): + """Test listing tokens with search and category filters.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + # Add test tokens + for token in TEST_TOKENS: + otp_generator.add_token( + token["secret"], + token["issuer"], + token["account"], + token["type"], + token["counter"], + category=token["category"], + tags=token["tags"] + ) + + # Test search by issuer + tokens = otp_generator.list_tokens(search="TestIssuer") + assert len(tokens) == 1 + assert tokens[0]["issuer"] == "TestIssuer" + + # Test search by tag + tokens = otp_generator.list_tokens(search="work") + assert len(tokens) == 1 + assert tokens[0]["category"] == "Work" + + # Test category filter + tokens = otp_generator.list_tokens(category="Personal") + assert len(tokens) == 1 + assert tokens[0]["category"] == "Personal" + +def test_get_categories(otp_generator): + """Test getting categories.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + # Add test tokens + for token in TEST_TOKENS: + otp_generator.add_token( + token["secret"], + token["issuer"], + token["account"], + token["type"], + token["counter"], + category=token["category"], + tags=token["tags"] + ) + + categories = otp_generator.get_categories() + assert len(categories) == 2 + assert "Work" in categories + assert "Personal" in categories + +def test_export_json(otp_generator): + """Test JSON export.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + # Add test tokens + for token in TEST_TOKENS: + otp_generator.add_token( + token["secret"], + token["issuer"], + token["account"], + token["type"], + token["counter"], + category=token["category"], + tags=token["tags"] + ) + + # Export to JSON + export_data = otp_generator.export_tokens("json") + tokens = json.loads(export_data) + assert len(tokens) == 2 + assert tokens[0]["issuer"] == TEST_TOKENS[0]["issuer"] + assert tokens[1]["issuer"] == TEST_TOKENS[1]["issuer"] + assert tokens[0]["category"] == TEST_TOKENS[0]["category"] + assert tokens[1]["category"] == TEST_TOKENS[1]["category"] + +def test_export_google(otp_generator): + """Test Google Authenticator export.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + # Add test tokens + for token in TEST_TOKENS: + otp_generator.add_token( + token["secret"], + token["issuer"], + token["account"], + token["type"], + token["counter"] + ) + + # Export to Google format + export_data = otp_generator.export_tokens("google") + lines = export_data.strip().split("\n") + assert len(lines) == 2 + assert "otpauth://" in lines[0] + assert "otpauth://" in lines[1] + assert "secret=" in lines[0] + assert "issuer=" in lines[0] + +def test_import_json(otp_generator): + """Test JSON import.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + # Create test JSON data + json_data = json.dumps(TEST_TOKENS) + + # Import from JSON + otp_generator.import_tokens(json_data, "json") + tokens = otp_generator.list_tokens() + assert len(tokens) == 2 + assert tokens[0]["issuer"] == TEST_TOKENS[0]["issuer"] + assert tokens[1]["issuer"] == TEST_TOKENS[1]["issuer"] + assert tokens[0]["category"] == TEST_TOKENS[0]["category"] + assert tokens[1]["category"] == TEST_TOKENS[1]["category"] + +def test_import_google(otp_generator): + """Test Google Authenticator import.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + # Create test Google format data + google_data = ( + "otpauth://totp/TestIssuer:test@example.com?secret=JBSWY3DPEHPK3PXP&issuer=TestIssuer\n" + "otpauth://hotp/TestIssuer2:test2@example.com?secret=JBSWY3DPEHPK3PXP&issuer=TestIssuer2&counter=0" + ) + + # Import from Google format + otp_generator.import_tokens(google_data, "google") + tokens = otp_generator.list_tokens() + assert len(tokens) == 2 + assert tokens[0]["issuer"] == "TestIssuer" + assert tokens[1]["issuer"] == "TestIssuer2" + assert tokens[0]["type"] == "totp" + assert tokens[1]["type"] == "hotp" + +def test_backup_restore(otp_generator): + """Test backup and restore functionality.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + # Add test tokens + for token in TEST_TOKENS: + otp_generator.add_token( + token["secret"], + token["issuer"], + token["account"], + token["type"], + token["counter"], + category=token["category"], + tags=token["tags"] + ) + + # Create backup + backup_file = otp_generator.create_backup(TEST_PASSWORD) + assert Path(backup_file).exists() + + # Create new instance and restore + with tempfile.TemporaryDirectory() as temp_dir: + new_otp = OTPGenerator(base_dir=temp_dir) + new_otp.restore_backup(backup_file, TEST_PASSWORD) + + # Verify restored tokens + new_otp._initialize_encryption(TEST_PASSWORD) + tokens = new_otp.list_tokens() + assert len(tokens) == 2 + assert tokens[0]["issuer"] == TEST_TOKENS[0]["issuer"] + assert tokens[1]["issuer"] == TEST_TOKENS[1]["issuer"] + assert tokens[0]["category"] == TEST_TOKENS[0]["category"] + assert tokens[1]["category"] == TEST_TOKENS[1]["category"] + +def test_qr_generation(otp_generator): + """Test QR code generation.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + # Add test token + otp_generator.add_token( + TEST_TOKENS[0]["secret"], + TEST_TOKENS[0]["issuer"], + TEST_TOKENS[0]["account"], + TEST_TOKENS[0]["type"] + ) + + # Generate QR code to file + with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_file: + otp_generator.generate_qr(0, temp_file.name) + assert Path(temp_file.name).exists() + assert Path(temp_file.name).stat().st_size > 0 + +def test_validate_tokens(otp_generator): + """Test token validation.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + # Add test tokens + for token in TEST_TOKENS: + otp_generator.add_token( + token["secret"], + token["issuer"], + token["account"], + token["type"], + token["counter"] + ) + + # Validate tokens + results = otp_generator.validate_tokens() + assert len(results) == 2 + assert all(result["is_valid"] for result in results) + + # Test with invalid token + otp_generator.add_token( + "invalid_secret", + "InvalidIssuer", + "invalid@example.com", + "totp" + ) + results = otp_generator.validate_tokens() + assert not all(result["is_valid"] for result in results) + assert any("Invalid secret" in result["error"] for result in results) + +def test_invalid_token_type(otp_generator): + """Test invalid token type handling.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + with pytest.raises(ValueError, match="Invalid token type"): + otp_generator.add_token( + TEST_TOKENS[0]["secret"], + TEST_TOKENS[0]["issuer"], + TEST_TOKENS[0]["account"], + "invalid_type" + ) + +def test_invalid_export_format(otp_generator): + """Test invalid export format handling.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + with pytest.raises(ValueError, match="Unsupported export format"): + otp_generator.export_tokens("invalid_format") + +def test_invalid_import_format(otp_generator): + """Test invalid import format handling.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + with pytest.raises(ValueError, match="Unsupported import format"): + otp_generator.import_tokens("{}", "invalid_format") + +def test_invalid_token_id(otp_generator): + """Test invalid token ID handling.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + with pytest.raises(ValueError, match="Invalid token ID"): + otp_generator.generate_otp(999) + +def test_encryption_initialization(otp_generator): + """Test encryption initialization.""" + with pytest.raises(ValueError, match="Not initialized"): + otp_generator.list_tokens() + +def test_backup_password_required(otp_generator): + """Test backup password requirement.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + with pytest.raises(ValueError, match="Password required"): + otp_generator.export_tokens("aegis") + +def test_password_strength(otp_generator): + """Test password strength checking.""" + # Test weak password + result = otp_generator.check_password_strength("weak") + assert not result["is_strong"] + assert result["score"] < 3 + + # Test strong password + result = otp_generator.check_password_strength("StrongP@ssw0rd123!") + assert result["is_strong"] + assert result["score"] >= 3 + +def test_clipboard_integration(otp_generator): + """Test clipboard integration.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + # Add test token + otp_generator.add_token( + TEST_TOKENS[0]["secret"], + TEST_TOKENS[0]["issuer"], + TEST_TOKENS[0]["account"], + TEST_TOKENS[0]["type"] + ) + + # Generate OTP and copy to clipboard + otp = otp_generator.generate_otp(0, copy_to_clipboard=True) + assert len(otp) == 6 + assert otp.isdigit() + +def test_token_metadata(otp_generator): + """Test token metadata handling.""" + otp_generator.install(TEST_PASSWORD) + otp_generator._initialize_encryption(TEST_PASSWORD) + + # Add token with metadata + otp_generator.add_token( + TEST_TOKENS[0]["secret"], + TEST_TOKENS[0]["issuer"], + TEST_TOKENS[0]["account"], + TEST_TOKENS[0]["type"], + category=TEST_TOKENS[0]["category"], + tags=TEST_TOKENS[0]["tags"] + ) + + # Verify metadata + tokens = otp_generator.list_tokens() + assert tokens[0]["category"] == TEST_TOKENS[0]["category"] + assert tokens[0]["tags"] == TEST_TOKENS[0]["tags"] + assert "created_at" in tokens[0] + assert tokens[0]["last_used"] is None + + # Generate OTP and check last_used + otp_generator.generate_otp(0) + tokens = otp_generator.list_tokens() + assert tokens[0]["last_used"] is not None \ No newline at end of file