#!/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, circle.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, circle.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 plus(self, x, y):
		xl = x-XOR_R
		xr = x+XOR_R
		yt = y-XOR_R
		yb = y+XOR_R
		size = XOR_R * 2
		self.t += f"""\
			<rect x="{xl}" y="{yt}" width="{size}" height="{size}"/>
			<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 minus(self, x, y):
		xl = x-XOR_R
		xr = x+XOR_R
		yt = y-XOR_R
		size = XOR_R * 2
		self.t += f"""\
			<rect x="{xl}" y="{yt}" width="{size}" height="{size}"/>
			<line x1="{xl}" y1="{y}" x2="{xr}" y2="{y}"/>
		"""
		return Block(self, x, y, XOR_R)

	def rotl(self, x, y, t):
		xl = x-XOR_R
		yt = y-XOR_R
		ymt = y-XOR_R/2
		ymb = y+XOR_R/2
		size = XOR_R * 2
		self.t += f"""\
			<rect x="{xl}" y="{yt}" width="{size}" height="{size}"/>
			<text class="t" x="{x}" y="{ymt}">&lt;&lt;&lt;</text>
			<text class="t" x="{x}" y="{ymb}">{t}</text>
		"""
		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)

	def block(self, x, y, wcells, hcells, texts, borders):
		cell_w = ENCRYPT_SIZE
		cell_h = ENCRYPT_SIZE//2
		w = cell_w * wcells
		h = cell_h * hcells
		xl = x-w//2
		yt = y-h//2
		self.t += f"""\
			<rect x="{xl}" y="{yt}" width="{w}" height="{h}"/>
		"""
		for (tx, ty, t) in texts:
			tx_ = (tx+0.5)*cell_w+xl
			ty_ = (ty+0.5)*cell_h+yt
			self.t += f"""<text class="t" x="{tx_}" y="{ty_}">{t}</text>\n"""
		return Block(self, x, y, ENCRYPT_SIZE//2)

	def knot(self, x, y):
		self.t += f"""\
			<circle class="f" r="4px" cx="{x}" cy="{y}"/>
		"""
		return Block(self, x, y, 0)

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)

# ChaCha Quarter Round
def chacha_qr():
	ax = 16 + 64*0
	bx = 16 + 64*1
	cx = 16 + 64*2
	dx = 16 + 64*3
	y = [64 + 32*i for i in range(12)]

	s = Svg()
	a1 = s.text(ax, 16, "a")
	b1 = s.text(bx, 16, "b")
	c1 = s.text(cx, 16, "c")
	d1 = s.text(dx, 16, "d")
	P1 = s.plus(ax, y[0])
	N1 = s.knot(bx, y[0])
	N2 = s.knot(ax, y[1])
	X1 = s.xor(dx, y[1])
	R1 = s.rotl(dx, y[2], "16")
	N3 = s.knot(dx, y[3])
	P2 = s.plus(cx, y[3])
	N4 = s.knot(cx, y[4])
	X2 = s.xor(bx, y[4])
	R2 = s.rotl(bx, y[5], "12")
	N5 = s.knot(bx, y[6])
	P3 = s.plus(ax, y[6])
	N6 = s.knot(ax, y[7])
	X3 = s.xor(dx, y[7])
	R3 = s.rotl(dx, y[8], "8")
	N7 = s.knot(dx, y[9])
	P4 = s.plus(cx, y[9])
	N8 = s.knot(cx, y[10])
	X4 = s.xor(bx, y[10])
	R4 = s.rotl(bx, y[11], "7")
	a2 = s.text(ax, y[-1]+48, "a")
	b2 = s.text(bx, y[-1]+48, "b")
	c2 = s.text(cx, y[-1]+48, "c")
	d2 = s.text(dx, y[-1]+48, "d")

	a1.to(P1)
	P1.to(P3)
	P3.to(a2)

	b1.to(X2)
	X2.to(R2)
	R2.to(X4)
	X4.to(R4)
	R4.to(b2)

	c1.to(P2)
	P2.to(P4)
	P4.to(c2)

	d1.to(X1)
	X1.to(R1)
	R1.to(X3)
	X3.to(R3)
	R3.to(d2)

	N1.to(P1)
	N2.to(X1)
	N3.to(P2)
	N4.to(X2)
	N5.to(P3)
	N6.to(X3)
	N7.to(P4)
	N8.to(X4)
	
	return SVG.format(body=s.t, title="ChaCha Quarter Round", w=372, h=y[-1]+64, **ARGS)

def chacha_encryption():
	s = Svg()
	B = s.block(64, 32, 4, 4, [
		(1.5, 0, "const"),
		(1.5, 1.5, "K"),
		(0.5, 3, "count"),
		(2.5, 3, "nonce"),
	],
	[])
	P = s.text(16, 144, "P")
	E = s.encrypt(64, 96)
	X = s.xor(64, 144)
	C = s.text(64, 180, "C")
	
	return SVG.format(body=s.t, title="ChaCha Encryption", 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())
	save("../flash-filesystem-encryption-2/chacha-encryption", chacha_encryption())