Symmetric file encryption using Argon2-based passphrases
import argparse
import secrets
import struct
import sys
from dataclasses import dataclass
from getpass import getpass
from pathlib import Path
from typing import List, NewType, Optional, Tuple

import argon2
from cryptography.hazmat.primitives import keywrap
from cryptography.hazmat.primitives.ciphers import aead

Salt = NewType("Salt", bytes)
Kek = NewType("Kek", bytes)

Dek = NewType("Dek", bytes)
WrappedDek = NewType("WrappedDek", bytes)

Nonce = NewType("Nonce", bytes)
CipherText = NewType("CipherText", bytes)
PlainText = NewType("PlainText", bytes)


@dataclass(frozen=True)
class Argon2Parameters:
    time_cost: int
    parallelism: int
    memory_cost: int


@dataclass(frozen=True)
class EncryptedFile:
    dek: WrappedDek
    salt: Salt
    nonce: Nonce
    argon2_parameters: Argon2Parameters
    cipher_text: CipherText

    @classmethod
    def read(cls, file: Path) -> "EncryptedFile":
        with file.open("rb") as reader:
            magic = reader.read(7)
            if magic != b"OBSCURA":
                raise ValueError("Unrecognized file format.")

            (
                dek,
                salt,
                nonce,
                argon2_time_cost,
                argon2_parallellism,
                argon2_memory_cost,
            ) = struct.unpack("40s32s12siii", reader.read(96))
            return cls(
                dek,
                salt,
                nonce,
                Argon2Parameters(
                    time_cost=argon2_time_cost,
                    parallelism=argon2_parallellism,
                    memory_cost=argon2_memory_cost,
                ),
                CipherText(reader.read()),
            )

    def to_bytes(self) -> bytes:
        return b"".join(
            [
                b"OBSCURA",
                self.dek,
                self.salt,
                self.nonce,
                struct.pack(
                    "iii",
                    self.argon2_parameters.time_cost,
                    self.argon2_parameters.parallelism,
                    self.argon2_parameters.memory_cost,
                ),
                self.cipher_text,
            ]
        )


def kdf(
    passphrase: bytes, *, salt: Optional[Salt] = None, parameters: Argon2Parameters
) -> Tuple[Kek, Salt]:
    if salt is None:
        salt = Salt(secrets.token_bytes(32))
    kek = Kek(
        argon2.low_level.hash_secret_raw(
            secret=passphrase,
            salt=salt,
            time_cost=parameters.time_cost,
            parallelism=parameters.parallelism,
            memory_cost=parameters.memory_cost,
            hash_len=32,
            type=argon2.Type.ID,
        )
    )
    return kek, salt


def generate_dek() -> Dek:
    return Dek(aead.AESGCM.generate_key(256))


def wrap_dek(kek: Kek, dek: Dek) -> WrappedDek:
    return WrappedDek(keywrap.aes_key_wrap(kek, dek))


def unwrap_dek(kek: Kek, dek: WrappedDek) -> Dek:
    return Dek(keywrap.aes_key_unwrap(kek, dek))


def encrypt(plain_text: PlainText, *, dek: Dek) -> Tuple[CipherText, Nonce]:
    nonce = Nonce(secrets.token_bytes(12))
    cipher_text = CipherText(aead.AESGCM(dek).encrypt(nonce, plain_text, None))
    return cipher_text, nonce


def decrypt(cipher_text: CipherText, *, dek: Dek, nonce: Nonce) -> PlainText:
    return PlainText(aead.AESGCM(dek).decrypt(nonce, cipher_text, None))


def encrypt_file(
    file: Path, *, passphrase: str, argon2_parameters: Argon2Parameters
) -> EncryptedFile:
    kek, salt = kdf(passphrase.encode(), parameters=argon2_parameters)
    dek = generate_dek()
    cipher_text, nonce = encrypt(PlainText(file.read_bytes()), dek=dek)
    return EncryptedFile(
        dek=wrap_dek(kek, dek),
        salt=salt,
        nonce=nonce,
        cipher_text=cipher_text,
        argon2_parameters=argon2_parameters,
    )


def decrypt_file(file: Path, *, passphrase: str) -> PlainText:
    enc_file = EncryptedFile.read(file)
    kek, _ = kdf(
        passphrase.encode(), salt=enc_file.salt, parameters=enc_file.argon2_parameters
    )
    return decrypt(
        cipher_text=enc_file.cipher_text,
        dek=unwrap_dek(kek, enc_file.dek),
        nonce=enc_file.nonce,
    )


def main(argv: List[str] = sys.argv[1:]) -> None:
    argparser = argparse.ArgumentParser()
    argparser.add_argument("-d", "--decrypt", action="store_true")
    argparser.add_argument(
        "--argon2-time-cost", type=int, default=argon2.DEFAULT_TIME_COST
    )
    argparser.add_argument(
        "--argon2-parallellism", type=int, default=argon2.DEFAULT_PARALLELISM
    )
    argparser.add_argument(
        "--argon2-memory-cost", type=int, default=argon2.DEFAULT_MEMORY_COST
    )
    argparser.add_argument("file", type=Path)
    args = argparser.parse_args(argv)

    argon2_parameters = Argon2Parameters(
        time_cost=args.argon2_time_cost,
        parallelism=args.argon2_parallellism,
        memory_cost=args.argon2_memory_cost,
    )

    passphrase = getpass()

    if args.decrypt:
        sys.stdout.buffer.write(decrypt_file(args.file, passphrase=passphrase))
    else:
        sys.stdout.buffer.write(
            encrypt_file(
                args.file,
                passphrase=passphrase,
                argon2_parameters=argon2_parameters,
            ).to_bytes(),
        )
    sys.stdout.flush()


if __name__ == "__main__":
    main()  # pragma: no cover

src/obscura.py
jocke-l/obscura
from pathlib import Path
from unittest import mock

import obscura
import pytest


@pytest.fixture
def argon2_parameters() -> obscura.Argon2Parameters:
    return obscura.Argon2Parameters(
        time_cost=1,
        parallelism=1,
        memory_cost=100,
    )


@pytest.fixture
def kek(argon2_parameters: obscura.Argon2Parameters) -> obscura.Kek:
    kek, _ = obscura.kdf(b"hunter2", parameters=argon2_parameters)
    return kek


@pytest.fixture
def dek() -> obscura.Dek:
    return obscura.generate_dek()


@pytest.fixture
def encrypted_file(
    argon2_parameters: obscura.Argon2Parameters,
) -> obscura.EncryptedFile:
    file = mock.Mock(spec=Path)
    file.read_bytes.return_value = b"secret snowflake"
    return obscura.encrypt_file(
        file, passphrase="hunter2", argon2_parameters=argon2_parameters  # noqa: S106
    )

tests/conftest.py
jocke-l/obscura
import secrets
from pathlib import Path
from unittest import mock

import obscura
import pytest


class TestKeys:
    def test_kek_length_equals_dek_length_of_32_bytes(
        self, kek: obscura.Kek, dek: obscura.Dek
    ) -> None:
        assert len(kek) == len(dek) == 32  # noqa: PLR2004

    def test_wrap_unwrap_identity(self, kek: obscura.Kek, dek: obscura.Dek) -> None:
        assert obscura.unwrap_dek(kek, obscura.wrap_dek(kek, dek)) == dek

    def test_wrapped_dek_is_different_than_unwrapped_dek(
        self, kek: obscura.Kek, dek: obscura.Dek
    ) -> None:
        assert obscura.wrap_dek(kek, dek) != obscura.WrappedDek(dek)


class TestEncryption:
    def test_encrypt_decrypt_identity(self, dek: obscura.Dek) -> None:
        plain_text = obscura.PlainText(b"foo")
        cipher_text, nonce = obscura.encrypt(plain_text, dek=dek)
        assert obscura.decrypt(cipher_text, dek=dek, nonce=nonce) == plain_text

    def test_decrypt_empty_cipher_text_raises_exception(self, dek: obscura.Dek) -> None:
        with pytest.raises(Exception):  # noqa: B017
            # TODO: Change this exception to somthing better
            obscura.decrypt(
                obscura.CipherText(b""),
                dek=dek,
                nonce=obscura.Nonce(secrets.token_bytes(12)),
            )


class TestEncryptedFile:
    def test_marshall_unmarshall_identity(
        self, encrypted_file: obscura.EncryptedFile
    ) -> None:
        file = mock.Mock(spec=Path)
        file.open = mock.mock_open(read_data=encrypted_file.to_bytes())
        assert obscura.EncryptedFile.read(file) == encrypted_file

    def test_unmarshall_invalid_file_format_raises_exception(self):
        file = mock.Mock(spec=Path)
        file.open = mock.mock_open(read_data=b"foobar")
        with pytest.raises(ValueError):
            # TODO: Change this exception to somtehing better
            obscura.EncryptedFile.read(file)

    def test_encrypt_decrypt_identity(
        self, encrypted_file: obscura.EncryptedFile
    ) -> None:
        with mock.patch("obscura.EncryptedFile.read") as read:
            read.return_value = encrypted_file
            file = mock.Mock(spec=Path)
            assert (
                obscura.decrypt_file(file, passphrase="hunter2")  # noqa: S106
                == b"secret snowflake"
            )


class TestCommandLine:
    def test_encrypt(self, capsysbinary: pytest.CaptureFixture[bytes]) -> None:
        getpass_patch = mock.patch("obscura.getpass", return_value="hunter2")
        read_bytes_patch = mock.patch(
            "obscura.Path.read_bytes", return_value=b"secret snowflake"
        )
        with getpass_patch, read_bytes_patch:
            obscura.main(["some_file"])

        captured = capsysbinary.readouterr()
        assert captured.out[:7] == b"OBSCURA"

    def test_decrypt(
        self,
        encrypted_file: obscura.EncryptedFile,
        capsysbinary: pytest.CaptureFixture[bytes],
    ) -> None:
        getpass_patch = mock.patch("obscura.getpass", return_value="hunter2")
        read_bytes_patch = mock.patch(
            "obscura.Path.open",
            mock.mock_open(read_data=encrypted_file.to_bytes()),
        )
        with getpass_patch, read_bytes_patch:
            obscura.main(["-d", "some_file"])

        captured = capsysbinary.readouterr()
        assert captured.out == b"secret snowflake"

tests/tests.py
jocke-l/obscura