#!/usr/bin/python3
"""
GNU AGPL v3, CopyLeft 2025 Pascal Engélibert [(why copyleft?)](https://txmn.tk/blog/why-copyleft/)
This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, version 3 of the License.  
This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details.  
You should have received a copy of the GNU Affero General Public License along with this program. If not, see https://www.gnu.org/licenses/.
"""

import math

ARGS = {
	"BLACK": "#000",
	"WHITE": "#ddd",
}
SVG = """\
<svg version="1.1" xmlns="http://www.w3.org/2000/svg" width="{w}" height="{h}" viewBox="0 0 {w} {h}">
	<!-- Generated by https://txmn.tk/blog/flash-filesystem-encryption/diagram.py -->
	<!-- Image released under license CC0 (public domain) -->
	<title>{title}</title>
	<style>
		text.t {{
			color: {BLACK};
			fill: {BLACK};
			text-anchor: middle;
			dominant-baseline: central;
			font-size: 16px;
			font-family: "Libertinus Sans";
		}}
		circle, line, rect, path.s {{
			stroke: {BLACK};
			stroke-width: 2px;
			fill: none;
		}}
		path.f {{
			stroke: none;
			fill: {BLACK};
		}}
		@media (prefers-color-scheme: dark) {{
			text.t {{
				color: {WHITE};
				fill: {WHITE};
			}}
			circle, line, rect, path.s {{
				stroke: {WHITE};
			}}
			path.f {{
				fill: {WHITE};
			}}
		}}
	</style>
{body}
</svg>
"""

class Block:
	def __init__(self, svg, x, y, r):
		self.svg = svg
		self.x = x
		self.y = y
		self.r = r

	def to(self, block):
		if self.x == block.x:
			if self.y > block.y:
				self.svg.t += arrow(self.x, self.y-self.r, block.x, block.y+block.r)
			else:
				self.svg.t += arrow(self.x, self.y+self.r, block.x, block.y-block.r)
		elif self.y == block.y:
			if self.x > block.x:
				self.svg.t += arrow(self.x-self.r, self.y, block.x+block.r, block.y)
			else:
				self.svg.t += arrow(self.x+self.r, self.y, block.x-block.r, block.y)

XOR_R = 8
ENCRYPT_SIZE = 32
class Svg:
	def __init__(self):
		self.t = ""

	def xor(self, x, y):
		xl = x-XOR_R
		xr = x+XOR_R
		yt = y-XOR_R
		yb = y+XOR_R
		self.t += f"""\
			<circle r="{XOR_R}" cx="{x}" cy="{y}"/>
			<line x1="{xl}" y1="{y}" x2="{xr}" y2="{y}"/>
			<line x1="{x}" y1="{yt}" x2="{x}" y2="{yb}"/>
		"""
		return Block(self, x, y, XOR_R)

	def encrypt(self, x, y):
		xl = x-ENCRYPT_SIZE//2
		yt = y-ENCRYPT_SIZE//2
		self.t += f"""\
			<rect x="{xl}" y="{yt}" width="{ENCRYPT_SIZE}" height="{ENCRYPT_SIZE}"/>
			<text class="t" x="{x}" y="{y}">E</text>
		"""
		return Block(self, x, y, ENCRYPT_SIZE//2)

	def text(self, x, y, t):
		self.t += f"""\
			<text class="t" x="{x}" y="{y}">{t}</text>
		"""
		return Block(self, x, y, 12)

	def square(self, x, y, t):
		xl = x-ENCRYPT_SIZE//2
		yt = y-ENCRYPT_SIZE//2
		self.t += f"""\
			<rect x="{xl}" y="{yt}" width="{ENCRYPT_SIZE}" height="{ENCRYPT_SIZE}"/>
			<text class="t" x="{x}" y="{y}">{t}</text>
		"""
		return Block(self, x, y, ENCRYPT_SIZE//2)

ARROW_TIP_START = 8
ARROW_TIP_BACK = 12
ARROW_TIP_R = 6
ARROW_TIP = [
	(-ARROW_TIP_START, 0),
	(-ARROW_TIP_BACK, ARROW_TIP_R),
	(0, 0),
	(-ARROW_TIP_BACK, -ARROW_TIP_R),
	(-ARROW_TIP_START, 0),
]
def arrow(x1, y1, x2, y2):
	angle = math.degrees(math.atan2(y2-y1, x2-x1))
	# Stop the line before the end to prevent it from passing through the tip
	length = math.sqrt((x2-x1)**2+(y2-y1)**2)
	xend = x2 - (x2-x1)/length*ARROW_TIP_START
	yend = y2 - (y2-y1)/length*ARROW_TIP_START
	path = ""
	px = 0
	py = 0
	for p in ARROW_TIP:
		path += "l{} {}".format(p[0]-px, p[1]-py)
		px = p[0]
		py = p[1]
	return f"""\
		<line x1="{x1}" y1="{y1}" x2="{xend}" y2="{yend}"/>
		<path class="f" transform="translate({x2},{y2}) rotate({angle})" d="M0 0{path}z"/>
	"""

def arrow_path(points):
	x1 = points[-2][0]
	y1 = points[-2][1]
	x2 = points[-1][0]
	y2 = points[-1][1]
	angle = math.degrees(math.atan2(y2-y1, x2-x1))
	# Stop the line before the end to prevent it from passing through the tip
	length = math.sqrt((x2-x1)**2+(y2-y1)**2)
	xend = x2 - (x2-x1)/length*ARROW_TIP_START
	yend = y2 - (y2-y1)/length*ARROW_TIP_START
	points[-1] = (xend, yend)
	path_arrow = ""
	px = 0
	py = 0
	for p in ARROW_TIP:
		path_arrow += "l{} {}".format(p[0]-px, p[1]-py)
		px = p[0]
		py = p[1]
	path_line = "M" + "L".join(["{} {}".format(p[0], p[1]) for p in points])
	return f"""\
		<path class="s" d="{path_line}"/>
		<path class="f" transform="translate({x2},{y2}) rotate({angle})" d="M0 0{path_arrow}z"/>
	"""

def ecb():
	s = Svg()
	for i in range(3):
		P = s.text(96*i+64, 16, "P")
		K = s.text(96*i+16, 64, "K")
		E = s.encrypt(96*i+64, 64)
		C = s.text(96*i+64, 112, "C")
		P.to(E)
		K.to(E)
		E.to(C)
	return SVG.format(body=s.t, title="ECB", w=288, h=128, **ARGS)

def ctr():
	s = Svg()
	for i in range(3):
		S = s.text(96*i+64, 16, f"N||{i}")
		K = s.text(96*i+16, 64, "K")
		E = s.encrypt(96*i+64, 64)
		P = s.text(96*i+16, 112, "P")
		X = s.xor(96*i+64, 112)
		C = s.text(96*i+64, 152, "C")
		S.to(E)
		K.to(E)
		E.to(X)
		P.to(X)
		X.to(C)
	return SVG.format(body=s.t, title="CTR", w=288, h=172, **ARGS)

def cbc():
	s = Svg()
	for i in range(3):
		P = s.text(96*i+64, 16, "P")
		X = s.xor(96*i+64, 56)
		K = s.text(96*i+16, 96, "K")
		E = s.encrypt(96*i+64, 96)
		C = s.text(96*i+64, 144, "C")
		P.to(X)
		X.to(E)
		K.to(E)
		E.to(C)
		if i == 0:
			IV = s.text(96*i+16, 56, "IV")
			IV.to(X)
		s.t += arrow_path([(96*i+72,144), (96*i+96,144), (96*i+96,56), (96*(1+i)+56,56)])
	s.text(356, 56, "···")
	return SVG.format(body=s.t, title="CBC", w=372, h=172, **ARGS)

def xts():
	s = Svg()
	P = s.text(224, 16, "P")
	X1 = s.xor(224, 56)
	E1 = s.encrypt(224, 96)
	X2 = s.xor(224, 144)
	C = s.text(224, 180, "C")

	K1 = s.text(176, 96, 'K<tspan dy="6">1</tspan>')

	I = s.text(64, 16, "i")
	E2 = s.encrypt(64, 96)

	K2 = s.text(16, 96, 'K<tspan dy="6">2</tspan>')

	J = s.text(128, 16, "j")
	A = s.square(128, 96, '×α <tspan dy="-8">j</tspan>')

	P.to(X1)
	X1.to(E1)
	K1.to(E1)
	I.to(E2)
	K2.to(E2)
	E1.to(X2)
	X2.to(C)
	E2.to(A)
	J.to(A)
	s.t += arrow_path([(144,96), (160,96), (160,56), (216,56)])
	s.t += arrow_path([(144,96), (160,96), (160,144), (216,144)])
	
	return SVG.format(body=s.t, title="XTS", w=372, h=192, **ARGS)

def save(name, data):
	f = open(f"{name}.svg", "w")
	f.write(data)
	f.close()

if __name__ == "__main__":
	save("ecb", ecb())
	save("ctr", ctr())
	save("cbc", cbc())
	save("xts", xts())