Skip to content

Commit f190964

Browse files
authored
Merge pull request #43
Implement streaming file encryption utilities
2 parents 064014c + a6896cb commit f190964

6 files changed

Lines changed: 155 additions & 10 deletions

File tree

ChatApp/file_crypto.py

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
from __future__ import annotations
2+
3+
import os
4+
from typing import Union
5+
6+
from cryptography.hazmat.backends import default_backend
7+
from cryptography.hazmat.primitives import hashes
8+
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
9+
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
10+
11+
12+
_CHUNK_SIZE = 4096
13+
_SALT_SIZE = 16
14+
_NONCE_SIZE = 12
15+
_TAG_SIZE = 16
16+
17+
18+
def _derive_key(password: Union[str, bytes], salt: bytes) -> bytes:
19+
"""Derive a 256-bit key from the given password and salt."""
20+
if isinstance(password, str):
21+
password = password.encode()
22+
kdf = PBKDF2HMAC(
23+
algorithm=hashes.SHA256(),
24+
length=32,
25+
salt=salt,
26+
iterations=390000,
27+
backend=default_backend(),
28+
)
29+
return kdf.derive(password)
30+
31+
32+
def encrypt_file(
33+
in_path: str,
34+
out_path: str,
35+
password: Union[str, bytes],
36+
chunk_size: int = _CHUNK_SIZE,
37+
) -> None:
38+
"""Encrypt ``in_path`` to ``out_path`` using AES-GCM.
39+
40+
The output file layout is ``salt`` + ``nonce`` + ciphertext + ``tag``. Data
41+
is processed in chunks so large files do not need to be fully loaded into
42+
memory.
43+
"""
44+
salt = os.urandom(_SALT_SIZE)
45+
key = _derive_key(password, salt)
46+
nonce = os.urandom(_NONCE_SIZE)
47+
cipher = Cipher(
48+
algorithms.AES(key),
49+
modes.GCM(nonce),
50+
backend=default_backend(),
51+
)
52+
encryptor = cipher.encryptor()
53+
54+
try:
55+
with open(in_path, "rb") as fin, open(out_path, "wb") as fout:
56+
fout.write(salt)
57+
fout.write(nonce)
58+
while True:
59+
chunk = fin.read(chunk_size)
60+
if not chunk:
61+
break
62+
data = encryptor.update(chunk)
63+
if data:
64+
fout.write(data)
65+
encryptor.finalize()
66+
fout.write(encryptor.tag)
67+
except Exception:
68+
if os.path.exists(out_path):
69+
os.remove(out_path)
70+
raise
71+
72+
73+
def decrypt_file(
74+
in_path: str,
75+
out_path: str,
76+
password: Union[str, bytes],
77+
chunk_size: int = _CHUNK_SIZE,
78+
) -> None:
79+
"""Decrypt ``in_path`` to ``out_path`` verifying the authentication tag."""
80+
total_size = os.path.getsize(in_path)
81+
if total_size < _SALT_SIZE + _NONCE_SIZE + _TAG_SIZE:
82+
raise ValueError("Ciphertext too small")
83+
84+
with open(in_path, "rb") as fin:
85+
salt = fin.read(_SALT_SIZE)
86+
nonce = fin.read(_NONCE_SIZE)
87+
key = _derive_key(password, salt)
88+
cipher = Cipher(
89+
algorithms.AES(key),
90+
modes.GCM(nonce),
91+
backend=default_backend(),
92+
)
93+
decryptor = cipher.decryptor()
94+
95+
ciphertext_len = total_size - _SALT_SIZE - _NONCE_SIZE - _TAG_SIZE
96+
remaining = ciphertext_len
97+
try:
98+
with open(out_path, "wb") as fout:
99+
while remaining > 0:
100+
chunk = fin.read(min(chunk_size, remaining))
101+
if not chunk:
102+
break
103+
remaining -= len(chunk)
104+
data = decryptor.update(chunk)
105+
if data:
106+
fout.write(data)
107+
tag = fin.read(_TAG_SIZE)
108+
decryptor.finalize_with_tag(tag)
109+
except Exception:
110+
if os.path.exists(out_path):
111+
os.remove(out_path)
112+
raise

Dockerfile

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
# Use an official Python runtime as a parent image
2-
FROM python:3.11-slim
2+
FROM python:3.11.13-slim-bullseye
33

44
LABEL maintainer="Psychevus"
55
LABEL description="Django WebSocket Chat App"
@@ -9,7 +9,6 @@ WORKDIR /app
99

1010
# Install system dependencies
1111
RUN apt-get update && \
12-
apt-get upgrade -y && \
1312
apt-get install -y --no-install-recommends gcc libffi-dev libssl-dev \
1413
default-libmysqlclient-dev && \
1514
rm -rf /var/lib/apt/lists/*

Dockerfile.demo

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
FROM python:3.11-slim
1+
FROM python:3.11.13-slim-bullseye
22

33
ENV PYTHONDONTWRITEBYTECODE=1 \
44
PYTHONUNBUFFERED=1

README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@ wscat -c ws://localhost:8000/ws/room/test/
8282

8383
```bash
8484
pip install -r requirements-dev.txt
85+
pip install -r requirements.txt # includes cryptography and pyOpenSSL for file encryption
8586
pytest --cov=ChatApp --cov=WebSocketChatApp
8687
```
8788

requirements.txt

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
11
Django~=4.2.5
2-
redis~=5.0.1
2+
redis~=6.2.0
33
channels~=3.0.5
44
django_ratelimit>=3.0.2
55
mysqlclient~=2.1.0
66
celery>=5.3,<6
77
django-allauth
8-
djangosaml2~=1.7.0
8+
djangosaml2~=1.11.1
99
django-scim2
10-
bleach~=6.1.0
10+
bleach~=6.2.0
1111
pyotp~=2.9.0
1212
kafka-python~=2.0.2
1313
PyJWT~=2.8.0
@@ -26,7 +26,9 @@ daphne>=3,<4
2626
python-json-logger~=3.3
2727
djangorestframework~=3.16.0
2828
drf-spectacular~=0.27.0
29-
pyfcm~=2.0
30-
apns2~=0.7
31-
32-
deprecated~=1.2.14
29+
pyfcm~=2.0
30+
apns2==0.7.1
31+
32+
deprecated~=1.2.14
33+
cryptography>=45.0,<46
34+
pyOpenSSL>=25.1,<26

tests/test_file_crypto.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
import os
2+
import tempfile
3+
4+
from ChatApp.file_crypto import encrypt_file, decrypt_file
5+
6+
7+
def _roundtrip(data: bytes, password: str = "pass"):
8+
with tempfile.TemporaryDirectory() as tmp:
9+
plain = os.path.join(tmp, "plain.bin")
10+
enc = os.path.join(tmp, "enc.bin")
11+
dec = os.path.join(tmp, "dec.bin")
12+
with open(plain, "wb") as fh:
13+
fh.write(data)
14+
encrypt_file(plain, enc, password)
15+
decrypt_file(enc, dec, password)
16+
with open(dec, "rb") as fh:
17+
return fh.read()
18+
19+
20+
def test_encrypt_decrypt_empty_file():
21+
assert _roundtrip(b"") == b""
22+
23+
24+
def test_encrypt_decrypt_small_file():
25+
data = b"hello world" * 5
26+
assert _roundtrip(data) == data
27+
28+
29+
def test_encrypt_decrypt_large_file():
30+
data = os.urandom(1024 * 1024) # 1MB
31+
assert _roundtrip(data) == data

0 commit comments

Comments
 (0)