#!/usr/bin/env python3

import sys
import os
import math
import struct
import subprocess
import tempfile
import getpass

HEADER_SIZE = 54
BITS_PER_PIXEL = 24
ITERATIONS = 200000

def bmp_geometry(size):
"""
Find a BMP geometry as close to a square as possible.
Each pixel row is aligned to a 4-byte boundary.
"""

target = math.sqrt(size * 4 / 3)
best = None

for width in range(max(1, int(target) - 2000),
int(target) + 2001):

stride = ((width * 3 + 3) // 4) * 4
height = (size + stride - 1) // stride
total = stride * height

score = abs(math.log(width / height)) + \
0.2 * (total - size) / size

if best is None or score < best[0]:
best = (score, width, height, stride, total)

return best[1:]

def main():

if len(sys.argv) != 3:
print(f"Usage: {sys.argv[0]} input_file output.bmp")
sys.exit(1)

source = sys.argv[1]
destination = sys.argv[2]
if not os.path.isfile(source):
sys.exit("Error: input file does not exist.")

password = getpass.getpass("Enter encryption password: ")
if not password:
sys.exit("Error: empty passwords are not allowed.")

# Encrypt the input file using OpenSSL.
print("Encrypting input file...")
with tempfile.NamedTemporaryFile() as tmp:
subprocess.run([
"openssl", "enc",
"-aes-256-cbc",
"-e",
"-nosalt",
"-pbkdf2",
"-iter", str(ITERATIONS),
"-pass", "stdin",
"-in", source,
"-out", tmp.name
], input=(password + "\n").encode(),
check=True)

tmp.seek(0)
ciphertext = tmp.read()

cipher_size = len(ciphertext)

# Calculate BMP dimensions.
width, height, stride, image_size = bmp_geometry(cipher_size)

# BMP file header (14 bytes).
file_header = struct.pack(
"<2sIHHI",
b"BM",
HEADER_SIZE + image_size,
0,
0,
HEADER_SIZE
)

# BITMAPINFOHEADER (40 bytes).
info_header = struct.pack(
"<IiiHHIIiiII",
40,
width,
height,
1,
BITS_PER_PIXEL,
0,
image_size,
2835,
2835,
0,
0
)

# Append zero padding to fill the BMP pixel array.
padding_size = image_size - cipher_size

with open(destination, "wb") as f:
f.write(file_header)
f.write(info_header)
f.write(ciphertext)
f.write(b"\x00" * padding_size)

print("\nBMP image successfully created.")
print(f"Output file : {destination}")
print(f"Image width : {width}")
print(f"Image height : {height}")
print(f"Ciphertext size : {cipher_size} bytes")
print(f"Padding size : {padding_size} bytes")
print(f"Ciphertext offset : {HEADER_SIZE}")

print("\nDecryption command:")
print(
f"dd if={destination} bs=1 skip=54 count={cipher_size} "
"| openssl enc -d -aes-256-cbc -nosalt -pbkdf2 "
f"-iter {ITERATIONS}"
)

print("\nThe password will be requested by OpenSSL during decryption.")

if __name__ == "__main__":
main()