(require 'seq)
(require 'hex-util)

(defmacro test (form)
  `(when (not ,form)
     (error "Assertion failed: %S" ',form)))

(defvar test-error-canary (make-symbol "test-error"))

(defmacro test-error (form)
  `(when (not (eq
               (condition-case nil ,form
                 (error test-error-canary))
               test-error-canary))
     (error "Expected error for %S" ',form)))

;; required for logcount and bignums
(test (version<= "27.1" emacs-version))

(defmacro test-equal (a b)
  `(when (not (equal ,a ,b))
     (error "Assertion failed: %S %S" ,a ,b)))

(defalias 'hex-to-bytes 'decode-hex-string)
(defalias 'base64-to-bytes 'base64-decode-string)

(defun bytesp (x)
  (and (stringp x) (not (multibyte-string-p x))))

(defun bytes-to-hex (bytes)
  (test (bytesp bytes))
  (encode-hex-string bytes))

(defun bytes-to-base64 (bytes)
  (test (bytesp bytes))
  (base64-encode-string bytes t))

(defun xor-bytes (a b)
  (test (bytesp a))
  (test (bytesp b))
  (test-equal (length a) (length b))
  (let ((bytes (make-string (length a) ?\s)))
    (dotimes (i (length bytes))
      (aset bytes i (logxor (aref a i) (aref b i))))
    (string-to-unibyte bytes)))

(defun xor-byte (bytes byte)
  (test (bytesp bytes))
  (test (and (>= byte 0) (< byte 256)))
  (let ((result (make-string (length bytes) ?\s)))
    (dotimes (i (length bytes))
      (aset result i (logxor (aref bytes i) byte)))
    (string-to-unibyte result)))

(defun xor-key (bytes key)
  (test (bytesp bytes))
  (test (bytesp key))
  (let ((result (make-string (length bytes) ?\s)))
    (dotimes (i (length bytes))
      (aset result i (logxor (aref bytes i) (aref key (% i (length key))))))
    (string-to-unibyte result)))

(defun file-contents (path)
  (with-temp-buffer
    (insert-file-contents path)
    (buffer-string)))

(defun file-lines (path)
  (split-string (file-contents path) "\n"))

(defvar english-histogram (make-hash-table))
(puthash ?\s 0.14 english-histogram)
(puthash ?e 0.12 english-histogram)
(puthash ?t 0.09 english-histogram)
(puthash 'other 0.09 english-histogram)
(puthash ?a 0.08 english-histogram)
(puthash ?o 0.07 english-histogram)
(puthash ?i 0.06 english-histogram)
(puthash ?n 0.06 english-histogram)
(puthash ?s 0.06 english-histogram)
(puthash ?h 0.06 english-histogram)
(puthash ?r 0.05 english-histogram)
(puthash ?d 0.04 english-histogram)
(puthash ?l 0.04 english-histogram)
(puthash ?u 0.02 english-histogram)
(puthash ?c 0.02 english-histogram)
(puthash ?m 0.02 english-histogram)
(puthash ?w 0.02 english-histogram)
(puthash ?f 0.02 english-histogram)
(puthash ?g 0.02 english-histogram)
(puthash ?y 0.01 english-histogram)
(puthash ?p 0.01 english-histogram)
(puthash ?b 0.01 english-histogram)
(puthash ?v 0.01 english-histogram)
(puthash ?k 0.01 english-histogram)
(puthash ?j 0.01 english-histogram)
(puthash ?x 0.00 english-histogram)
(puthash ?q 0.00 english-histogram)
(puthash ?z 0.00 english-histogram)

(defun frequencies (bytes)
  (let ((result (make-hash-table))
        (total (length bytes)))
    (dotimes (i total)
      (let ((char (aref bytes i)))
        (puthash char (1+ (gethash char result 0)) result)))
    (maphash (lambda (k v) (puthash k (/ (float v) total) result)) result)
    result))

(defun chi-squared (hist1 hist2)
  (let ((score 0))
    (maphash
     (lambda (k v1)
       (let ((v2 (gethash k hist2 0)))
         (when (> v1 0)
           (setq score (+ score (/ (expt (- v1 v2) 2) v1))))))
     hist1)
    score))

(defun printablep (string)
  (string-match-p "^[[:print:]]*$" string))

(defun english-score (bytes)
  (if (not (printablep bytes))
      0
    (setq bytes (downcase bytes))
    (setq bytes (replace-regexp-in-string "[^ a-z]" "." bytes))
    (let ((hist (frequencies bytes)))
      (puthash 'other (gethash ?. hist 0) hist)
      (remhash ?. hist)
      (let ((score (/ 1 (chi-squared english-histogram hist))))
        (if (< (gethash 'other hist) 0.05)
            (* score 2)
          score)))))

(defun hamming-distance (a b)
  (test (bytesp a))
  (test (bytesp b))
  (test-equal (length a) (length b))
  (let ((combination (xor-bytes a b))
        (distance 0))
    (dotimes (i (length combination))
      (setq distance (+ distance (logcount (aref combination i)))))
    distance))

(defun transpose-bytes (items)
  (test (consp items))
  (let ((length (length (car items))))
    (let (result)
      (dotimes (i length)
        (push
         (mapconcat (lambda (item)
                      (if (< i (length item)) (string (aref item i)) ""))
                    items
                    "")
         result))
      (nreverse result))))

(defvar aes-128-block-size 16)
(defvar aes-128-key-size 16)
(defvar aes-128-iv-size 16)
(defvar aes-128-null-iv (make-string aes-128-iv-size 0))

(defvar aes-128-cbc-id
  (plist-get (cdr (assq 'AES-128-CBC (gnutls-ciphers))) :cipher-id))

(defun aes-128-ecb-decrypt-block (block key)
  (setq key (copy-sequence key)) ; key is wiped
  (car (gnutls-symmetric-decrypt aes-128-cbc-id key aes-128-null-iv block)))
(defun aes-128-ecb-encrypt-block (block key)
  (setq key (copy-sequence key)) ; key is wiped
  (car (gnutls-symmetric-encrypt aes-128-cbc-id key aes-128-null-iv block)))

(defun aes-128-ecb-decrypt (ciphertext key)
  (test (bytesp ciphertext))
  (test (bytesp key))
  (test (zerop (% (length ciphertext) aes-128-block-size)))
  (test (= (length key) aes-128-block-size))
  (mapconcat (lambda (block) (aes-128-ecb-decrypt-block block key))
             (seq-partition ciphertext aes-128-block-size)
             ""))

(defun aes-128-ecb-encrypt (plaintext key)
  (test (bytesp plaintext))
  (test (bytesp key))
  (test (zerop (% (length plaintext) aes-128-block-size)))
  (test (= (length key) aes-128-block-size))
  (mapconcat (lambda (block) (aes-128-ecb-encrypt-block block key))
             (seq-partition plaintext aes-128-block-size)
             ""))

(defun repeated-block-p (blocks)
  (/= (length blocks) (length (seq-uniq blocks))))

(defun pkcs7pad (bytes &optional block-size)
  (test (bytesp bytes))
  (when (not block-size)
    (setq block-size aes-128-block-size))
  (let ((padsize (- block-size (% (length bytes) block-size))))
    (string-as-unibyte (concat bytes (make-string padsize padsize)))))

(defun pkcs7unpad (bytes &optional block-size)
  (test (bytesp bytes))
  (test (zerop (% (length bytes) aes-128-block-size)))
  (when (not block-size)
    (setq block-size aes-128-block-size))
  (let ((padsize (aref bytes (1- (length bytes)))))
    (test (<= padsize block-size))
    (let ((padding (substring bytes (- (length bytes) padsize))))
      (test (seq-every-p (lambda (c) (= c padsize)) padding))
      (substring bytes 0 (- (length bytes) padsize)))))

(defun aes-128-cbc-decrypt (ciphertext key iv)
  (test (bytesp ciphertext))
  (test (bytesp key))
  (test (bytesp iv))
  (test (zerop (% (length ciphertext) aes-128-block-size)))
  (test (= (length key) aes-128-block-size))
  (test (= (length iv) aes-128-block-size))
  (let ((blocks (seq-partition ciphertext aes-128-block-size))
        result)
    (dolist (block blocks)
      (push (xor-bytes (aes-128-ecb-decrypt-block block key) iv) result)
      (setq iv block))
    (mapconcat 'identity (nreverse result) "")))

(defun aes-128-cbc-encrypt (plaintext key iv)
  (test (bytesp plaintext))
  (test (bytesp key))
  (test (bytesp iv))
  (test (zerop (% (length plaintext) aes-128-block-size)))
  (test (= (length key) aes-128-block-size))
  (test (= (length iv) aes-128-block-size))
  (let ((blocks (seq-partition plaintext aes-128-block-size))
        result)
    (dolist (block blocks)
      (setq iv (aes-128-ecb-encrypt-block (xor-bytes block iv) key))
      (push iv result))
    (mapconcat 'identity (nreverse result) "")))

(defun random-integer (from to)
  (+ (random (- to from)) from))

(defun random-bytes (length &optional to)
  (let ((bytes (make-string length ?\s)))
    (dotimes (i (length bytes))
      (aset bytes i (random-integer 0 (or to 256))))
    (string-to-unibyte bytes)))