elsie-four/elsie-four.fth

\ Elsie-Four (LC4), Copyright 2018 etb, License: GPLv3+
CREATE K 36 ALLOT           \ Key buffer
CREATE S 36 ALLOT           \ State buffer
CREATE M 0 ,                \ Marker

: ->INT ( c -- n)
    DUP                  [CHAR] #  =     IF DROP 0             ELSE
    DUP                  [CHAR] _  =     IF DROP 1             ELSE
    DUP [CHAR] 2 >= OVER [CHAR] 9 <= AND IF [ CHAR 2 2  - ]L - ELSE
    DUP [CHAR] a >= OVER [CHAR] z <= AND IF [ CHAR a 10 - ]L - ELSE
    DUP [CHAR] A >= OVER [CHAR] Z <= AND IF [ CHAR A 10 - ]L - ELSE
    ." Warning: mapping '" EMIT ." ' to '_'" CR  1
    THEN THEN THEN THEN THEN ; 
: ->CHAR ( n -- c)  >R
    S" #_23456789abcdefghijklmnopqrstuvwxyz" R@ <
    IF [CHAR] * ELSE R@ + C@ THEN  R> DROP ; 

CREATE BUF 256 ALLOT  CREATE B BUF ,
: B0   BUF B ! ;   : .B   BUF B @ BUF - TYPE ;
: B,   B @ C!  [ 1 CHARS ]L B +! ;

CREATE A 0 ,
: C@A+ ( -- c)  A @ C@  1 CHARS A +! ;
: A-C! ( c --)  A @ 1 CHARS - TUCK C!  A ! ;
: RIGHT-ROTATE ( row --) 6 * S +  A !
    C@A+ C@A+ C@A+ C@A+ C@A+ C@A+ >R
    A-C! A-C! A-C! A-C! A-C! R> A-C! ;
: C@A6+ ( -- c) A @ C@  6 CHARS A +! ;
: A6-C! ( c --) A @ 6 CHARS - TUCK C!  A ! ;
: DOWN-ROTATE ( col --) S +  A !
    C@A6+ C@A6+ C@A6+ C@A6+ C@A6+ C@A6+ >R
    A6-C! A6-C! A6-C! A6-C! A6-C! R> A6-C! ;

: S[] ( n -- c)  S + C@ ;
: SFIND ( c -- n)  >R
    0 BEGIN DUP S[] R@ <> WHILE CHAR+ REPEAT R> DROP ;
: +S ( n n' -- n'')      \ Add indices within the state matrix
    6 /MOD ROT  6 /MOD ROT
    + 6 MOD 6 * >R  + 6 MOD  R> + ;
: -S ( n n' -- n'')      \ Subtract indices within the state matrix
    6 /MOD ROT  6 /MOD ROT
    - 6 MOD 6 * >R  SWAP - 6 MOD  R> + ;

\ Rather than maintain row and column indices for various markers,
\ just keep track of the character, and search, via SFIND, for the
\ index in S when needed.
: UPDATE ( C P --)
        SFIND  6 /  RIGHT-ROTATE    \ rotate row of P
    DUP SFIND  6 MOD DOWN-ROTATE    \ rotate column of C
    M @ SFIND +S S[] M ! ;          \ adjust marker
: CIPHER ( c --) ->INT
    DUP SFIND  M @ ( P P' M)
    +S ( P C')  S[] TUCK ( C P C)
    ->CHAR B,  UPDATE ;
: PLAIN ( c --) ->INT
    DUP SFIND  M @ ( C C' M)
    -S ( C P')  S[] DUP ( C P P)
    ->CHAR B,  UPDATE ;

: (ENCRYPT) ( c-addr u)  BOUNDS  ?DO I C@ CIPHER LOOP ;
: (DECRYPT) ( c-addr u)  BOUNDS  ?DO I C@ PLAIN  LOOP ;

: RESET    K S 36 CMOVE  0 S[] M ! ;   \ Reset the state matrix and marker
: ENCRYPT ( nonce u1 header u2 plaintext u3 sig u4)
    RESET
    2>R 2>R 2SWAP      \ save plaintext/sig for later, setup nonce
    (ENCRYPT) B0       \ encrypt the nonce and ignore
    (ENCRYPT) B0       \ encrypt header, if any, and ignore
    2R> 2R> 2SWAP      \ restore sig/plaintext
    (ENCRYPT)          \ encrypt the plaintext
    (ENCRYPT)          \ append the encrypted sig
    CR ." Ciphertext: " .B CR ;
: DECRYPT ( nonce u1 header u2 ciphertext u3)
    RESET
    2SWAP 2ROT         \ save ciphertext for later, setup nonce
    (ENCRYPT) B0       \ encrypt the nonce and ignore
    (ENCRYPT) B0       \ encrypt header, if any, and ignore
    (DECRYPT)          \ decrypt the ciphertext
    CR ." Plaintext: " .B CR ;

: S.  S A ! 6 0 DO 6 0 DO C@A+ ->CHAR EMIT LOOP SPACE LOOP ;   
: M.  SPACE M @ SFIND 6 /MOD . SPACE . SPACE ;
: TRACE ( c-addr u) CR  RESET
    ." State" 38 SPACES ." i  j  pt ct" CR
    S. M. CR  BOUNDS ?DO 
    I C@ CIPHER
    S. M. I C@ EMIT 2 SPACES B @ 1 CHARS - C@ EMIT CR
    LOOP ;

\ Convenience syntax
: SEND
    BL PARSE                \ nonce
    BL PARSE                \ header
    BL PARSE                \ plaintext
    BL PARSE                \ sig
    ENCRYPT ;
: RECEIVE
    BL PARSE                \ nonce
    BL PARSE                \ header
    BL PARSE                \ ciphertext
    DECRYPT ;

: PRIV-KEY ( " ccc" --)
    BL WORD  COUNT
    36 <> IF DROP ." ERROR: Key must have length 36" CR QUIT THEN
    A !  36 0 DO C@A+ ->INT  K I + C! LOOP ; IMMEDIATE
PRIV-KEY XV7YDQ#OPAJ_39RZUT8B45WCSGEHMIKNF26L
RESET