(load (expand-file-name "../util" (file-name-directory (or load-file-name buffer-file-name))))
(require 'seq)

(setq plaintext (base64-to-bytes "Um9sbGluJyBpbiBteSA1LjAKV2l0aCBteSByYWctdG9wIGRvd24gc28gbXkgaGFpciBjYW4gYmxvdwpUaGUgZ2lybGllcyBvbiBzdGFuZGJ5IHdhdmluZyBqdXN0IHRvIHNheSBoaQpEaWQgeW91IHN0b3A/IE5vLCBJIGp1c3QgZHJvdmUgYnkK"))
(setq key (random-bytes aes-128-key-size))
(setq random-prefix (random-bytes (random-integer 5 10)))
(message "Prefixsize: %d" (length random-prefix))

(defun encryption-oracle (input)
  (aes-128-ecb-encrypt (pkcs7pad (concat random-prefix input plaintext)) key))

(defun detect-blocksize ()
  (let ((input "")
        (last-length (length (encryption-oracle ""))))
    (catch 'exit
      (while t
        (setq input (concat input "A"))
        (let ((new-length (length (encryption-oracle input))))
          (when (> new-length last-length)
            (throw 'exit (- new-length last-length)))
          (setq last-length new-length))))))

(defun detect-ecb (blocksize)
  (let ((plaintext (make-string (* blocksize 3) ?A)))
    (repeated-block-p (seq-partition (encryption-oracle plaintext) blocksize))))

;; |XXXXXXXXXXAAAAAB|BBBBBBBBBBBBBBBB|BBBBBBBBBBBBBBBP|
;; |XXXXXXXXXXAAAAAA|BBBBBBBBBBBBBBBB|BBBBBBBBBBBBBBBB|PPPPPPPPPPPPPPPP|
(defun guess-prefixsize (blocksize)
  (let ((prefix ""))
    (catch 'exit
      (while t
        (let* ((input (concat prefix (make-string (* blocksize 2) ?B)))
               (blocks (seq-partition (encryption-oracle input) blocksize)))
          (when (repeated-block-p blocks)
            (throw 'exit (- blocksize (length prefix))))
          (setq prefix (concat prefix "A")))))))

;; |XXXXXXXXXXAAAAA?|????????????????|??????PPPPPPPPPP| 10+5+0=15
;; |XXXXXXXXXXAAAAs?|????????????????|?????PPPPPPPPPPP| 10+4+1=15
;; |XXXXXXXXXXAAAse?|????????????????|????PPPPPPPPPPPP| 10+3+2=15
;; |XXXXXXXXXXAAsec?|????????????????|???PPPPPPPPPPPPP| 10+2+3=15
;; |XXXXXXXXXXAsecr?|????????????????|??PPPPPPPPPPPPPP| 10+1+4=15
;; |XXXXXXXXXXsecre?|????????????????|?PPPPPPPPPPPPPPP| 10+0+5=15
;; |XXXXXXXXXXAAAAAA|AAAAAAAAAsecret?|????????????????|PPPPPPPPPPPPPPPP| 10+15+6=31
;; |XXXXXXXXXXAAAAAA|AAAAAAAAsecret1?|???????????????P| 10+14+7=31
;; |XXXXXXXXXXAAAAAA|AAAAAAAsecret12?|??????????????PP| 10+13+8=31
;; |XXXXXXXXXXAAAAAA|AAAAAAsecret123?|?????????????PPP| 10+12+9=31
;; |XXXXXXXXXXAAAAAA|AAAAAsecret123s?|????????????PPPP| 10+11+10=31
;; |XXXXXXXXXXAAAAAA|AAAAsecret123se?|???????????PPPPP| 10+10+11=31
;; |XXXXXXXXXXAAAAAA|AAAsecret123sec?|??????????PPPPPP| 10+9+12=31
;; |XXXXXXXXXXAAAAAA|AAsecret123secr?|?????????PPPPPPP| 10+8+13=31
;; |XXXXXXXXXXAAAAAA|Asecret123secre?|????????PPPPPPPP| 10+7+14=31
;; |XXXXXXXXXXAAAAAA|secret123secret?|???????PPPPPPPPP| 10+6+15=31
;; |XXXXXXXXXXAAAAAs|ecret123secret4?|??????PPPPPPPPPP| 10+5+16=31
;; |XXXXXXXXXXAAAAse|cret123secret45?|?????PPPPPPPPPPP| 10+4+17=31
;; |XXXXXXXXXXAAAsec|ret123secret456?|????PPPPPPPPPPPP| 10+3+18=31
;; |XXXXXXXXXXAAsecr|et123secret456s?|???PPPPPPPPPPPPP| 10+2+19=31
;; |XXXXXXXXXXAsecre|t123secret456se?|??PPPPPPPPPPPPPP| 10+1+20=31
;; |XXXXXXXXXXsecret|123secret456sec?|?PPPPPPPPPPPPPPP| 10+0+21=31
;; |XXXXXXXXXXAAAAAA|AAAAAAAAAsecret1|23secret456secr?|PPPPPPPPPPPPPPPP| 10+15+22=47
;; |XXXXXXXXXXAAAAAA|AAAAAAAAsecret12|3secret456secreP| 10+14+23=47
(defun guess-byte (known prefixsize blocksize)
  (let* ((short-block-count (1+ (/ (+ (length known) prefixsize) blocksize)))
         (padsize (- (* short-block-count blocksize) (length known) prefixsize 1))
         (input (make-string padsize ?A))
         (block-idx (/ (+ prefixsize padsize (length known)) blocksize))
         (block (elt (seq-partition (encryption-oracle input) blocksize)
                     block-idx)))
    (catch 'return
      (dotimes (i 128)
        (let* ((input (concat (make-string padsize ?A) known (string i)))
               (match (elt (seq-partition (encryption-oracle input) blocksize)
                           block-idx)))
          (when (equal block match)
            (throw 'return i))))
      (throw 'return nil))))

(setq blocksize (detect-blocksize))
(message "Detected blocksize: %d" blocksize)
(when (detect-ecb blocksize)
  (message "Detected ECB"))
(setq prefixsize (guess-prefixsize blocksize))
(message "Detected prefixsize: %d" prefixsize)
(let ((known "")
      guessed-bytes
      guessed-byte)
  (while (setq guessed-byte (guess-byte known prefixsize blocksize))
    (princ (format "%c" guessed-byte))
    (send-string-to-terminal "") ; flush stdout
    (setq known (concat known (string guessed-byte)))))