#!/usr/bin/env python3

import base64
import hmac
import os
import re
import stat
import struct
import subprocess
import sys
import time

# Names are printed to the terminal; control characters are excluded.
NAME_RE = re.compile(r"[A-Za-z0-9._@+-]+")
DIGESTS = ("sha1", "sha256", "sha512")
DEFAULT_SECRETS = "~/.config/otp_secrets"


class SecretsError(Exception):
    def __init__(self, message, status=1):
        super().__init__(message)
        self.status = status


def decode_key(key):
    if not key:
        raise ValueError("secret must not be empty")
    try:
        decoded = base64.b32decode(key.upper() + "=" * ((8 - len(key)) % 8))
    except ValueError as error:
        raise ValueError("secret is not valid base32") from error
    if not decoded:
        raise ValueError("secret decodes to an empty key")
    return decoded


def validate(period, digits, digest):
    if not 1 <= period <= 3600:
        raise ValueError("period must be between 1 and 3600")
    if not 6 <= digits <= 10:
        raise ValueError("digits must be between 6 and 10")
    if digest not in DIGESTS:
        raise ValueError(f"digest must be one of {', '.join(DIGESTS)}")


def hotp(key, counter, digits=6, digest="sha1"):
    mac = hmac.new(decode_key(key), struct.pack(">Q", counter), digest).digest()
    offset = mac[-1] & 0x0F
    binary = struct.unpack(">L", mac[offset : offset + 4])[0] & 0x7FFFFFFF
    return str(binary)[-digits:].zfill(digits)


def totp(key, period=30, digits=6, digest="sha1"):
    return hotp(key, int(time.time()) // period, digits, digest)


def secrets_text():
    cmd = os.environ.get("OTP_SECRETS_CMD")
    if cmd:
        proc = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE)
        if proc.returncode != 0:
            raise SecretsError("secrets command failed", proc.returncode)
        output = proc.stdout
        source = "OTP_SECRETS_CMD"
    else:
        path = os.environ.get("OTP_SECRETS") or os.path.expanduser(DEFAULT_SECRETS)
        try:
            with open(path, "rb") as f:
                info = os.fstat(f.fileno())
                if not stat.S_ISREG(info.st_mode):
                    raise SecretsError(f"{path}: not a regular file")
                if info.st_uid != os.geteuid():
                    raise SecretsError(f"{path}: not owned by the current user")
                if info.st_mode & 0o077:
                    raise SecretsError(f"{path}: permissions must be 0600 or stricter")
                output = f.read()
        except OSError as error:
            raise SecretsError(f"{path}: {error.strerror}") from None
        source = path
    try:
        return output.decode(), source
    except UnicodeDecodeError:
        raise SecretsError(f"{source}: not valid UTF-8") from None


def parse_entries(text, source):
    entries = []
    names = set()
    for lineno, line in enumerate(text.splitlines(), 1):
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        fields = line.split()
        if not 2 <= len(fields) <= 5:
            raise SecretsError(f"{source}:{lineno}: expected 2-5 fields")
        name, key = fields[0], fields[1]
        if not NAME_RE.fullmatch(name):
            raise SecretsError(f"{source}:{lineno}: invalid name {name!r}")
        if name in names:
            raise SecretsError(f"{source}:{lineno}: duplicate name {name!r}")
        try:
            period = int(fields[2]) if len(fields) > 2 else 30
            digits = int(fields[3]) if len(fields) > 3 else 6
        except ValueError:
            raise SecretsError(
                f"{source}:{lineno}: period and digits must be integers"
            ) from None
        digest = fields[4].lower() if len(fields) > 4 else "sha1"
        try:
            validate(period, digits, digest)
            decode_key(key)
        except ValueError as error:
            raise SecretsError(f"{source}:{lineno}: {error}") from None
        entries.append((name, key, period, digits, digest))
        names.add(name)
    return entries


def load_entries():
    text, source = secrets_text()
    return parse_entries(text, source)


def main():
    if len(sys.argv) > 4:
        sys.exit("usage: otp.py [period [digits [digest]]]")
    try:
        period = int(sys.argv[1]) if len(sys.argv) > 1 else 30
        digits = int(sys.argv[2]) if len(sys.argv) > 2 else 6
        digest = sys.argv[3].lower() if len(sys.argv) > 3 else "sha1"
        validate(period, digits, digest)
        for key in sys.stdin:
            print(totp(key.strip(), period, digits, digest))
    except ValueError as error:
        sys.exit(f"otp: {error}")


if __name__ == "__main__":
    main()