# Neural AES (NNAES)

A PyTorch implementation of the AES cryptographic algorithm using **Deep Neural Networks (DNNs)**, based on the concept of *Deep Neural Cryptography*.

---

## Attribution

This project is **adapted from the original work**:

- Repository: https://github.com/DavidGerault/deep_neural_cryptography  
- Paper: *Deep Neural Cryptography* – Gerault et al.  
  https://eprint.iacr.org/2025/288.pdf  

The original implementation was reorganized and extended for usability, packaging, and experimentation purposes.

GNU General Public License v3 (GPLv3)

---

## Description

This library implements AES-128 as a **deterministic neural network**, without any training process.

The model operates on binary inputs represented as tensors of shape:
```
(batch_size, 128)
```

Each AES block is encoded as a 128-dimensional binary vector.

### Key characteristics

- Exact implementation of AES operations using neural networks  
- No learning or training required (fully deterministic)  
- Based on:
  - Linear layers + ReLU activations  
  - Corner functions for Boolean logic  
  - Fixed, analytically constructed weights  

### Input / Output utilities

- `integer_to_bitvector`: converts a 128-bit integer into a binary tensor  
- `state_as_ints`: converts model outputs back to integers  

---

## What was done in this project

Compared to the original implementation, this project introduces:

- A clean and reusable **Python package structure**  
- Improved modularity and readability  
- Integration with **PyTorch workflows**  
- Benchmarking tools for performance analysis  
- Unit tests based on AES standard test vectors (FIPS / AESAVS)  

---

## Installation

Clone the repository and install locally:

```bash
pip install -e .
```
For development (tests, etc.):
```bash
pip install -e .[dev]
```
---
## Usage

### Basic example (encryption protected)
```python
from nnaes import NeuralAES
from nnaes.utils import integer_to_bitvector, state_as_ints
import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

key = 0x2b7e151628aed2a6abf7158809cf4f3c
plaintext = 0x3243f6a8885a308d313198a2e0370734

model = NeuralAES(
    secret_key=key,
    direction="Encryption",
    protected=True,
    epsilon=1/4
).to(device)

model.eval()

x = torch.tensor(
    [integer_to_bitvector(plaintext)],
    dtype=torch.float16
).to(device)

with torch.inference_mode():
    y = model(x)

ciphertext = state_as_ints(y)[0]
print(f"{ciphertext:032x}")
```

Expected output: 

```
3925841d02dc09fbdc118597196a0b32
```
---
## Notes
- The implementation uses torch.float16 by default for performance reasons.
- GPU execution is recommended for efficient benchmarking.
- On CPU, using float32 may improve numerical stability.
---
## Disclaimer

This project is intended for research and educational purposes only.
It is not optimized for production cryptographic use.