use std::{
	cmp::{min},
	ffi::{OsStr},
	io::{self, Read, Write, stdin, stdout},
	process::{exit}
};

const ROUNDS: usize = 128;

struct StreamCipher {
	ks: KeySchedule,
	counter: [u128; 2],
	buf: [u128; 2],
	index: usize
}

struct KeySchedule {
	ks: [u128; ROUNDS]
}

impl KeySchedule {
	fn new(key: &[u128; 2]) -> Self {
		let mut ks = [0; ROUNDS];

		let mut y = key[0];
		let mut x = key[1];

		for i in 0..ROUNDS {
			Self::round(&mut x, &mut y, i as _);
			ks[i] = y;
		}

		KeySchedule { ks }
	}

	fn round(x: &mut u128, y: &mut u128, k: u128) {
		*x = x.rotate_right(8).wrapping_add(*y) ^ k;
		*y = y.rotate_left(3) ^ *x;
	}

	fn encrypt(&self, msg: &[u128; 2]) -> [u128; 2] {
		let mut y = msg[0];
		let mut x = msg[1];

		for &k in self.ks.iter() {
			Self::round(&mut x, &mut y, k);
		}

		[y, x]
	}
}

impl StreamCipher {
	fn new(key: &[u128; 2], nonce: u128) -> Self {
		let ks = KeySchedule::new(key);
		let counter = [0, nonce];

		StreamCipher {
			buf: ks.encrypt(&counter),
			index: 0,
			ks,
			counter
		}
	}

	fn buf(&self) -> &[u8; 32] {
		unsafe { &*(self.buf.as_ptr() as *const [u8; 32]) }
	}

	fn crypt(&mut self, mut buf: &mut [u8]) {
		while !buf.is_empty() {
			if self.index >= self.buf.len() {
				self.counter[0] = self.counter[0].checked_add(1).unwrap();
				self.buf = self.ks.encrypt(&self.counter);
				self.buf[0] = self.buf[0].to_le();
				self.buf[1] = self.buf[1].to_le();
				self.index = 0;
			}

			let len = min(buf.len(), self.buf()[self.index..].len());

			for (a, &b) in buf[..len].iter_mut().zip(&self.buf()[self.index..][..len]) {
				*a ^= b;
			}

			buf = &mut { buf }[len..];
			self.index += len;
		}
	}
}

fn main() {
	let mut args = std::env::args_os();
	let name = args.next().unwrap_or_default();

	if args.len() != 2 {
		eprintln!("{} <nonce> <key>", name.to_string_lossy());
		exit(1);
	}

	let mut stream_cipher = {
		let nonce = parse_hex(&args.next().unwrap(), "nonce", 32);
		assert!(nonce[1] == 0);
		let key = parse_hex(&args.next().unwrap(), "key", 64);

		StreamCipher::new(&key, nonce[0])
	};

	let stdin = stdin();
	let mut stdin = stdin.lock();

	let stdout = stdout();
	let mut stdout = stdout.lock();

	let mut buf = vec!(0; 8 << 10).into_boxed_slice();

	loop {
		match stdin.read(&mut buf) {
			Ok(0) => { break; }
			Ok(read) => {
				stream_cipher.crypt(&mut buf[..read]);

				match stdout.write_all(&buf[..read]) {
					Ok(()) => {}
					Err(ref e) if e.kind() == io::ErrorKind::BrokenPipe => { break; }
					Err(e) => {
						eprintln!("{}", e);
						exit(1);
					}
				}
			}
			Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
			Err(e) => {
				eprintln!("{}", e);
				exit(1);
			}
		}
	}
}

fn parse_hex(s: &OsStr, name: &str, max_len: usize) -> [u128; 2] {
	fn val(b: u8, name: &str) -> u8 {
		match b {
			b'0'...b'9' => { b - b'0' }
			b'a'...b'f' => { b - b'a' + 10 }
			b'A'...b'F' => { b - b'A' + 10 }
			_ => {
				eprintln!("{} must be hexadecimal", name);
				exit(1);
			}
		}
	}

	let bytes = s.to_str().unwrap().as_bytes();

	if bytes.len() > max_len {
		eprintln!("{} must not be longer than {} hex characters", name, max_len);
		exit(1);
	}

	let mut ret = [0; 2];

	{
		let ret = unsafe { &mut *(ret.as_mut_ptr() as *mut [u8; 32]) };

		for (c, r) in bytes.chunks(2).zip(ret) {
			*r = (val(c[0], name) << 4) | val(c[1], name);
		}
	}

	[u128::from_le(ret[0]), u128::from_le(ret[1])]
}