p.ari.lt/tests/test.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Test API"""

import base64
import hashlib
import hmac
import os
import random
import secrets
import sys
import time
from typing import Any, Dict, Optional
from warnings import filterwarnings as filter_warnings

import requests
import sympy
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305


def chacha20_encrypt(data: bytes, subkey: bytes) -> str:
    """Encrypt data using a subkey"""
    nonce: bytes = os.urandom(12)
    key: bytes = hashlib.sha3_256(
        nonce + subkey[:48] + int(time.time() / 256).to_bytes(8, "big")
    ).digest()
    return base64.b85encode(
        nonce + ChaCha20Poly1305(key).encrypt(nonce, data, subkey[48:])
    ).decode("ascii")


def chacha20_decrypt(ct: str, subkey: bytes) -> bytes:
    """Decrypt data using a subkey"""
    cb: bytes = base64.b85decode(ct.encode("ascii"))
    key: bytes = hashlib.sha3_256(
        cb[:12] + subkey[:48] + int(time.time() / 256).to_bytes(8, "big")
    ).digest()
    return ChaCha20Poly1305(key).decrypt(cb[:12], cb[12:], subkey[48:])


def generate_large_prime(n: int) -> int:
    """Generate a large n-bit prime number"""

    found_prime: bool = False
    p: int = 0

    while not found_prime:
        p = secrets.randbits(n)  # Generates a random n-bit number
        p |= (
            1 << n - 1
        ) | 1  # Ensures the n-bit length and assures it is an odd number
        found_prime = (
            sympy.isprime(p)
            and sympy.isprime((2 * p) + 1)
            and sympy.isprime((p - 1) // 2)
        )  # Generate safe primes

    return p


def hmac_sha3_512(key: bytes, message: bytes) -> bytes:
    """Generate HMAC using SHA3-512."""
    return hmac.new(key, message, hashlib.sha3_512).digest()


def hotp(key: bytes, counter: int) -> int:
    """Generate a HOTP code based on the key and counter value."""

    counter_bytes: bytes = counter.to_bytes(8, byteorder="big")
    hmac_digest: bytes = hmac_sha3_512(key, counter_bytes)

    offset: int = hmac_digest[-1] & 0x0F
    code: bytes = hmac_digest[offset : offset + 4]

    code_int: int = int.from_bytes(code, byteorder="big") & 0x7FFFFFFF

    return code_int % (10**10)


def totp(key: bytes, interval: int = 256, reference_ts: Optional[int] = None) -> int:
    """Generate a TOTP code."""

    counter: int = (
        reference_ts if reference_ts is not None else int(time.time())
    ) // interval

    return hotp(key, counter)


def validate_totp(
    provided_totp: int,
    key: bytes,
    interval: int = 256,
    reference_ts: Optional[int] = None,
) -> bool:
    """Validate the provided TOTP code."""

    current_time: int = reference_ts if reference_ts is not None else int(time.time())

    for offset in [-1, 0, 1]:  # Check the previous, current, and next intervals
        if provided_totp == totp(key, interval, current_time + offset * interval):
            return True

    return False


def main() -> int:
    """entry / main function"""

    tkey: Dict[str, Any] = requests.get("http://127.0.0.1:8080/api/totp").json()
    tkey["key"] = base64.b85decode(tkey["key"])

    print(
        requests.get(
            f"http://127.0.0.1:8080/api/ranges?totp={totp(tkey[     'key'], reference_ts=tkey['ts'])}"
        ).json()
    )

    if len(sys.argv) < 2:
        p: int = generate_large_prime(128)
    else:
        p: int = int(sys.argv[1])

    print(len(bin(p)) - 2)
    print(
        requests.get(
            f"http://127.0.0.1:8080/api/requirements/{len(bin(p))-1}?totp={totp(tkey['key'], reference_ts=tkey['ts'])}",
        ).json()
    )

    addprime: Dict[str, Any] = requests.post(
        "http://127.0.0.1:8080/api/add",
        json={
            "prime": str(p),
            "totp": totp(tkey["key"], reference_ts=tkey["ts"]),
        },
    ).json()

    addprime["subkey"] = base64.b85decode(addprime["subkey"].encode("ascii"))
    addprime["key"] = chacha20_decrypt(addprime["key"], addprime["subkey"])
    addprime["totp_key"] = chacha20_decrypt(addprime["totp_key"], addprime["subkey"])

    note: Dict[str, Any] = requests.post(
        "http://127.0.0.1:8080/api/note",
        json={
            "prime": str(p),
            "key": chacha20_encrypt(addprime["key"], addprime["subkey"]),
            "note": f"This is a note for {p}.",
            "totp": totp(addprime["totp_key"], reference_ts=addprime["totp_ts"]),
        },
    ).json()

    print(note)

    addprime: Dict[str, Any] = requests.post(
        "http://127.0.0.1:8080/api/rekey",
        json={
            "prime": str(p),
            "totp": totp(addprime["totp_key"], reference_ts=addprime["totp_ts"]),
            "key": chacha20_encrypt(addprime["key"], addprime["subkey"]),
        },
    ).json()

    print(addprime)

    addprime["subkey"] = base64.b85decode(addprime["subkey"].encode("ascii"))
    addprime["key"] = chacha20_decrypt(addprime["key"], addprime["subkey"])
    addprime["totp_key"] = chacha20_decrypt(addprime["totp_key"], addprime["subkey"])

    note: Dict[str, Any] = requests.post(
        "http://127.0.0.1:8080/api/note",
        json={
            "prime": str(p),
            "key": chacha20_encrypt(addprime["key"], addprime["subkey"]),
            "note": f"This is a REKEYED note for {p}.",
            "totp": totp(addprime["totp_key"], reference_ts=addprime["totp_ts"]),
        },
    ).json()

    getprime: Dict[str, Any] = requests.get(
        f"http://127.0.0.1:8080/api/prime/{p}?totp={totp(tkey[     'key'], reference_ts=tkey['ts'])}"
    ).json()

    print(getprime)

    return 0


if __name__ == "__main__":
    assert main.__annotations__.get("return") is int, "main() should return an integer"

    filter_warnings("error", category=Warning)
    raise SystemExit(main())