(**************************************************************************)
(*                                                                        *)
(*  Encryption library                                                    *)
(*  Copyright (C) 2017   Peter Moylan                                     *)
(*                                                                        *)
(*  This program is free software: you can redistribute it and/or modify  *)
(*  it under the terms of the GNU General Public License as published by  *)
(*  the Free Software Foundation, either version 3 of the License, or     *)
(*  (at your option) any later version.                                   *)
(*                                                                        *)
(*  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 General Public License for more details.                          *)
(*                                                                        *)
(*  You should have received a copy of the GNU General Public License     *)
(*  along with this program.  If not, see <http://www.gnu.org/licenses/>. *)
(*                                                                        *)
(*  To contact author:   http://www.pmoylan.org   peter@pmoylan.org       *)
(*                                                                        *)
(**************************************************************************)


<* WOFF316+ *>

IMPLEMENTATION MODULE AES;

        (********************************************************)
        (*                                                      *)
        (*    AES: Keyed hashing for message authentication     *)
        (*                                                      *)
        (*  Implementation of AES as defined in NIST.FIPS.197   *)
        (*                                                      *)
        (*  Three variants are defined by the standard:         *)
        (*      AES-128     4-word key, 10 rounds               *)
        (*      AES-192     6-word key, 12 rounds               *)
        (*      AES-256     8-word key, 14 rounds               *)
        (*                                                      *)
        (*  The CBC mode is defined in RFC 3602.                *)
        (*  CTR mode is defined in RFC 3686, but RFC4344 uses   *)
        (*  a different definition for use with SSH.  Because   *)
        (*  of this conflict, we define a new mode SDCTR,       *)
        (*  initiated by a call to AES_SDCTRopen, that          *)
        (*  differs from CTR mode by following RFC4344.         *)
        (*                                                      *)
        (*  Programmer:         P. Moylan                       *)
        (*  Started:            12 October 2017                 *)
        (*  Last edited:        4 November 2021                 *)
        (*  Status:          Passes all tests in NIST.FIPS.197  *)
        (*                                                      *)
        (*      For CBC, passes cases 1..4 in RFC 3602.  I have *)
        (*      not yet had the patience to run the other cases.*)
        (*                                                      *)
        (*      For CTR, passes every test in RFC 3686.         *)
        (*                                                      *)
        (*      FOR SDCTR, passes tests in NIST.FIPS.800-38a.   *)
        (*                                                      *)
        (*      Suspicion of final padding error, retest this.  *)
        (*                                                      *)
        (********************************************************)

FROM SYSTEM IMPORT
    (* type *)  CARD8, CARD32,
    (* proc *)  CAST, ADR;

FROM Storage IMPORT
    (* proc *)  ALLOCATE, DEALLOCATE;

FROM LowLevel IMPORT
    (* proc *)  IXOR, IXORB, Copy, AddOffset;

(************************************************************************)

CONST
    Nb = 4;       (* number of columns in the state *)

TYPE
    (* The context record holds the Nk value and the round keys.  Since *)
    (* there is not a lot of difference in size of the round keys for   *)
    (* the different variants, we choose to allocate space for the      *)
    (* biggest case rather than have a variable-sized record.           *)

    (* If incount = 16 then the "inbuff" array contains a copy of the   *)
    (* input that was processed in the previous operation.  Otherwise   *)
    (* that array is partially filled with incount bytes of input       *)
    (* that was not processed in the previous operation because there   *)
    (* was not enough to fill a complete block.   The "state" array     *)
    (* holds the result of the last encryption or decryption.           *)

    (* The counter array is used only in CTR mode.                      *)
    (* The blockwaiting flag says that there is a decrypted result      *)
    (* in the state array that has not yet been sent to the output.     *)
    (* We need to put a one-block delay in decryption if the full       *)
    (* padding option is in use, because we don't know whether the      *)
    (* block contains padding until we get to the very last block.      *)
    (* This, however, is only when fullpadding is specified.  That      *)
    (* option is not enabled in the CTR case.                           *)

    AEScontext = POINTER TO
                        RECORD
                            decrypt, CBC, CTR, SDCTR,
                                    fullpadding, blockwaiting: BOOLEAN;
                            incount: CARDINAL;
                            inbuff: StateArray;
                            counter: StateArray;
                            state: StateArray;
                            Nk: CARDINAL;
                            roundkeys: ARRAY [0..14] OF StateArray;
                        END (*RECORD*);

    FourByte = ARRAY [0..3] OF CARD8;

    (* The state in the AES algorithm is a 4x4 array of bytes, but we   *)
    (* find it convenient to store each column as a single 32-bit word. *)
    (* For the purposes of copying between input and state, and state   *)
    (* and output, the array is scanned in column major ordering.       *)
    (* This arrangement does require that to access A[row,col], we      *)
    (* have to express it as A[col].byte[row], but that is a minor      *)
    (* detail.                                                          *)

    (* Remark: the input block, the output block, and the state all     *)
    (* have the same size.                                              *)

    Word =  RECORD
                CASE :BOOLEAN OF
                    FALSE:  word: CARD32;
                    |
                    TRUE:   byte: FourByte;
                END (*CASE*);
            END (*RECORD*);

    StateArray = ARRAY [0..Nb-1] OF Word;

    (* An Sbox is a byte-to-byte mapping used to implement some of the  *)
    (* steps in the calculation.  The standard specifies it as a 16x16  *)
    (* array, but it is more efficient to implement it as a straight    *)
    (* CARD8 to CARD8 mapping.                                          *)

    Sbox = ARRAY CARD8 OF CARD8;

(************************************************************************)

TYPE RconType = ARRAY [1..10] OF CARDINAL;

CONST
    (* Rcon is a constant array used in key expansion. *)

    Rcon = RconType {01H, 02H, 04H, 08H, 10H, 20H, 40H, 80H, 1BH, 36H};

(************************************************************************)
(*                            THE S-BOXES                               *)
(************************************************************************)

CONST
    Sboxdir = Sbox{
                 063H, 07CH, 077H, 07BH, 0F2H, 06BH, 06FH, 0C5H, 030H, 001H, 067H, 02BH, 0FEH, 0D7H, 0ABH, 076H,
                 0CAH, 082H, 0C9H, 07DH, 0FAH, 059H, 047H, 0F0H, 0ADH, 0D4H, 0A2H, 0AFH, 09CH, 0A4H, 072H, 0C0H,
                 0B7H, 0FDH, 093H, 026H, 036H, 03FH, 0F7H, 0CCH, 034H, 0A5H, 0E5H, 0F1H, 071H, 0D8H, 031H, 015H,
                 004H, 0C7H, 023H, 0C3H, 018H, 096H, 005H, 09AH, 007H, 012H, 080H, 0E2H, 0EBH, 027H, 0B2H, 075H,
                 009H, 083H, 02CH, 01AH, 01BH, 06EH, 05AH, 0A0H, 052H, 03BH, 0D6H, 0B3H, 029H, 0E3H, 02FH, 084H,
                 053H, 0D1H, 000H, 0EDH, 020H, 0FCH, 0B1H, 05BH, 06AH, 0CBH, 0BEH, 039H, 04AH, 04CH, 058H, 0CFH,
                 0D0H, 0EFH, 0AAH, 0FBH, 043H, 04DH, 033H, 085H, 045H, 0F9H, 002H, 07FH, 050H, 03CH, 09FH, 0A8H,
                 051H, 0A3H, 040H, 08FH, 092H, 09DH, 038H, 0F5H, 0BCH, 0B6H, 0DAH, 021H, 010H, 0FFH, 0F3H, 0D2H,
                 0CDH, 00CH, 013H, 0ECH, 05FH, 097H, 044H, 017H, 0C4H, 0A7H, 07EH, 03DH, 064H, 05DH, 019H, 073H,
                 060H, 081H, 04FH, 0DCH, 022H, 02AH, 090H, 088H, 046H, 0EEH, 0B8H, 014H, 0DEH, 05EH, 00BH, 0DBH,
                 0E0H, 032H, 03AH, 00AH, 049H, 006H, 024H, 05CH, 0C2H, 0D3H, 0ACH, 062H, 091H, 095H, 0E4H, 079H,
                 0E7H, 0C8H, 037H, 06DH, 08DH, 0D5H, 04EH, 0A9H, 06CH, 056H, 0F4H, 0EAH, 065H, 07AH, 0AEH, 008H,
                 0BAH, 078H, 025H, 02EH, 01CH, 0A6H, 0B4H, 0C6H, 0E8H, 0DDH, 074H, 01FH, 04BH, 0BDH, 08BH, 08AH,
                 070H, 03EH, 0B5H, 066H, 048H, 003H, 0F6H, 00EH, 061H, 035H, 057H, 0B9H, 086H, 0C1H, 01DH, 09EH,
                 0E1H, 0F8H, 098H, 011H, 069H, 0D9H, 08EH, 094H, 09BH, 01EH, 087H, 0E9H, 0CEH, 055H, 028H, 0DFH,
                 08CH, 0A1H, 089H, 00DH, 0BFH, 0E6H, 042H, 068H, 041H, 099H, 02DH, 00FH, 0B0H, 054H, 0BBH, 016H};

    Sboxinv = Sbox{
                 052H, 009H, 06AH, 0D5H, 030H, 036H, 0A5H, 038H, 0BFH, 040H, 0A3H, 09EH, 081H, 0F3H, 0D7H, 0FBH,
                 07CH, 0E3H, 039H, 082H, 09BH, 02FH, 0FFH, 087H, 034H, 08EH, 043H, 044H, 0C4H, 0DEH, 0E9H, 0CBH,
                 054H, 07BH, 094H, 032H, 0A6H, 0C2H, 023H, 03DH, 0EEH, 04CH, 095H, 00BH, 042H, 0FAH, 0C3H, 04EH,
                 008H, 02EH, 0A1H, 066H, 028H, 0D9H, 024H, 0B2H, 076H, 05BH, 0A2H, 049H, 06DH, 08BH, 0D1H, 025H,
                 072H, 0F8H, 0F6H, 064H, 086H, 068H, 098H, 016H, 0D4H, 0A4H, 05CH, 0CCH, 05DH, 065H, 0B6H, 092H,
                 06CH, 070H, 048H, 050H, 0FDH, 0EDH, 0B9H, 0DAH, 05EH, 015H, 046H, 057H, 0A7H, 08DH, 09DH, 084H,
                 090H, 0D8H, 0ABH, 000H, 08CH, 0BCH, 0D3H, 00AH, 0F7H, 0E4H, 058H, 005H, 0B8H, 0B3H, 045H, 006H,
                 0D0H, 02CH, 01EH, 08FH, 0CAH, 03FH, 00FH, 002H, 0C1H, 0AFH, 0BDH, 003H, 001H, 013H, 08AH, 06BH,
                 03AH, 091H, 011H, 041H, 04FH, 067H, 0DCH, 0EAH, 097H, 0F2H, 0CFH, 0CEH, 0F0H, 0B4H, 0E6H, 073H,
                 096H, 0ACH, 074H, 022H, 0E7H, 0ADH, 035H, 085H, 0E2H, 0F9H, 037H, 0E8H, 01CH, 075H, 0DFH, 06EH,
                 047H, 0F1H, 01AH, 071H, 01DH, 029H, 0C5H, 089H, 06FH, 0B7H, 062H, 00EH, 0AAH, 018H, 0BEH, 01BH,
                 0FCH, 056H, 03EH, 04BH, 0C6H, 0D2H, 079H, 020H, 09AH, 0DBH, 0C0H, 0FEH, 078H, 0CDH, 05AH, 0F4H,
                 01FH, 0DDH, 0A8H, 033H, 088H, 007H, 0C7H, 031H, 0B1H, 012H, 010H, 059H, 027H, 080H, 0ECH, 05FH,
                 060H, 051H, 07FH, 0A9H, 019H, 0B5H, 04AH, 00DH, 02DH, 0E5H, 07AH, 09FH, 093H, 0C9H, 09CH, 0EFH,
                 0A0H, 0E0H, 03BH, 04DH, 0AEH, 02AH, 0F5H, 0B0H, 0C8H, 0EBH, 0BBH, 03CH, 083H, 053H, 099H, 061H,
                 017H, 02BH, 004H, 07EH, 0BAH, 077H, 0D6H, 026H, 0E1H, 069H, 014H, 063H, 055H, 021H, 00CH, 07DH};

(************************************************************************)
(*                     MULTIPLICATION IN GF(2^8)                        *)
(************************************************************************)

PROCEDURE xtime (val: CARD8): CARD8;

    (* Multiples the polynomial represented by val by x, as described   *)
    (* in section 4.2.1 of the standard.                                *)

    BEGIN
        IF val >= 128 THEN
            (* High-order bit is set. *)
            RETURN IXORB (1BH, 2*(val-128));
        ELSE
            RETURN 2*val;
        END (*IF*);
    END xtime;

(************************************************************************)

PROCEDURE Prod (M, N: CARD8): CARD8;

    (* Product of M and N, as described in section 4.2.1 of the standard. *)

    VAR power, result: CARD8;

    BEGIN
        result := 0;  power := M;
        WHILE N <> 0 DO
            IF ODD(N) THEN
                result := IXORB(result, power);
            END (*IF*);
            power := xtime(power);
            N := N DIV 2;
        END (*WHILE*);
        RETURN result;
    END Prod;

(************************************************************************)
(*                     TRANSFORMATION OPERATIONS                        *)
(************************************************************************)

PROCEDURE BlockXOR (VAR (*IN*) mask: StateArray;
                            VAR (*INOUT*) result: StateArray);

    (* XOR addition of mask to result. *)

    VAR j: CARDINAL;

    BEGIN
        FOR j := 0 TO 3 DO
            result[j].word := IXOR (result[j].word, mask[j].word);
        END (*FOR*);
    END BlockXOR;

(************************************************************************)

PROCEDURE AddRoundKey (VAR (*INOUT*) state: StateArray;
                            VAR (*IN*) key: StateArray);

    (* XOR addition of key to state. *)

    BEGIN
        BlockXOR (key, state);
    END AddRoundKey;

(************************************************************************)

PROCEDURE InvMixColumns (VAR (*INOUT*) state: StateArray);

    (* Inverse of MixColumns, see below.                *)
    (* Mixes data in the state to produce new columns.  *)

    VAR T: FourByte;
        j: CARDINAL;

    BEGIN
        FOR j := 0 TO 3 DO
            (* Let T be the original column j of state. *)
            T := state[j].byte;
            state[j].byte[0] := IXORB (IXORB(Prod(T[0], 0EH), Prod(T[1],0BH)),
                                       IXORB(Prod(T[2], 0DH), Prod (T[3], 09H)));
            state[j].byte[1] := IXORB (IXORB(Prod(T[0], 09H), Prod(T[1],0EH)),
                                       IXORB(Prod(T[2], 0BH), Prod (T[3], 0DH)));
            state[j].byte[2] := IXORB (IXORB(Prod(T[0], 0DH), Prod(T[1],09H)),
                                       IXORB(Prod(T[2], 0EH), Prod (T[3], 0BH)));
            state[j].byte[3] := IXORB (IXORB(Prod(T[0], 0BH), Prod(T[1],0DH)),
                                       IXORB(Prod(T[2], 09H), Prod (T[3], 0EH)));
        END (*FOR*);
    END InvMixColumns;

(************************************************************************)

PROCEDURE InvShiftRows (VAR (*INOUT*) state: StateArray);

    (* Inverse of ShiftRows, see below.  Cyclic shifting of the rows of *)
    (* the state by different offsets.                                  *)

    VAR i, j, col: CARDINAL;
        temp: StateArray;

    BEGIN
        temp := state;
        FOR i := 1 TO 3 DO
            FOR j := 0 TO 3 DO
                col := (j + i) MOD 4;
                state[col].byte[i] := temp[j].byte[i];
            END (*FOR*);
        END (*FOR*);
    END InvShiftRows;

(************************************************************************)

PROCEDURE InvSubBytes (VAR (*INOUT*) state: StateArray);

    (* Inverse of SubBytes, see below.  Transformation in the cipher    *)
    (* that processes the state using a nonlinear byte substitution     *)
    (* table (S-box) that operates on each of the state bytes           *)
    (* independently.                                                   *)

    VAR i, j: CARDINAL;

    BEGIN
        FOR j := 0 TO 3 DO
            FOR i := 0 TO 3 DO
                state[j].byte[i] := Sboxinv[state[j].byte[i]];
            END (*FOR*);
        END (*FOR*);
    END InvSubBytes;

(************************************************************************)

PROCEDURE MixColumns (VAR (*INOUT*) state: StateArray);

    (* Mixes data in the state to produce new columns. *)

    VAR T: FourByte;
        j: CARDINAL;

    BEGIN
        FOR j := 0 TO 3 DO
            (* Let T be the original column j of state. *)
            T := state[j].byte;
            state[j].byte[0] := IXORB (IXORB(Prod(T[0], 2), Prod(T[1],3)), IXORB(T[2], T[3]));
            state[j].byte[1] := IXORB (IXORB(T[0], Prod(T[1], 2)), IXORB(Prod(T[2], 3), T[3]));
            state[j].byte[2] := IXORB (IXORB(T[0], T[1]), IXORB(Prod(T[2], 2), Prod(T[3], 3)));
            state[j].byte[3] := IXORB (IXORB(Prod(T[0], 3), T[1]), IXORB(T[2], Prod(T[3], 2)));
        END (*FOR*);
    END MixColumns;

(************************************************************************)

PROCEDURE RotWord (A: FourByte): FourByte;

    (* Cyclic permutation of a four-byte word. *)

    VAR i: CARDINAL;  result: FourByte;

    BEGIN
        FOR i := 0 TO 2 DO
            result[i] := A[i+1];
        END (*FOR*);
        result[3] := A[0];
        RETURN result;
    END RotWord;

(************************************************************************)

PROCEDURE ShiftRows (VAR (*INOUT*) state: StateArray);

    (* Cyclic shifting of the rows of the state by different offsets. *)

    VAR i, j, col: CARDINAL;
        temp: StateArray;

    BEGIN
        temp := state;
        FOR i := 1 TO 3 DO
            FOR j := 0 TO 3 DO
                col := (j + i) MOD 4;
                state[j].byte[i] := temp[col].byte[i];
            END (*FOR*);
        END (*FOR*);
    END ShiftRows;

(************************************************************************)

PROCEDURE SubWord (A: FourByte): FourByte;

    (* Function used in the Key Expansion routine that takes a          *)
    (* four-byte input word and applies an S-box to each of the four    *)
    (* bytes to produce an output word.                                 *)

    VAR i: CARDINAL;  result: FourByte;

    BEGIN
        FOR i := 0 TO 3 DO
            result[i] := Sboxdir[A[i]];
        END (*FOR*);
        RETURN result;
    END SubWord;

(************************************************************************)

PROCEDURE SubBytes (VAR (*INOUT*) state: StateArray);

    (* Transformation in the cipher that processes the state using a    *)
    (* nonlinear byte substitution table (S-box) that operates on each  *)
    (* of the state bytes independently.                                *)

    VAR i, j: CARDINAL;

    BEGIN
        FOR j := 0 TO 3 DO
            FOR i := 0 TO 3 DO
                state[j].byte[i] := Sboxdir[state[j].byte[i]];
            END (*FOR*);
        END (*FOR*);
    END SubBytes;

(************************************************************************)
(*                 TRANSFORMATION OF ONE BLOCK OF DATA                  *)
(************************************************************************)

PROCEDURE Cipher (Nr: CARDINAL;  VAR (*IN*) indata: StateArray;
                                VAR (*OUT*) outdata: StateArray;
                                VAR (*IN*) w: ARRAY OF StateArray);

    (* The main encryption algorithm. *)

    VAR state: StateArray;
        round: CARDINAL;

    BEGIN
        state := indata;
        AddRoundKey(state, w[0]);           (* See Sec. 5.1.4   *)
        FOR round := 1 TO Nr-1 DO
            SubBytes(state);                (* See Sec. 5.1.1   *)
            ShiftRows(state);               (* See Sec. 5.1.2   *)
            MixColumns(state);              (* See Sec. 5.1.3   *)
            AddRoundKey(state, w[round]);
        END (*FOR*);
        SubBytes(state);
        ShiftRows(state);
        AddRoundKey(state, w[Nr]);
        outdata := state;
    END Cipher;

(************************************************************************)

PROCEDURE InvCipher (Nr: CARDINAL;  VAR (*IN*) in: StateArray;
                                    VAR (*OUT*) out: StateArray;
                                    VAR (*IN*) w: ARRAY OF StateArray);

    (* The main decryption algorithm.  Note that an alternative method  *)
    (* is given in section 5.3.5 of the standard, which is supposed to  *)
    (* be more efficient because it uses the same sequence of steps as  *)
    (* the encryption algorithm.  I have carefully examined this, and   *)
    (* cannot see any gain in efficiency, so I have not implemented the *)
    (* alternative algorithm.                                           *)

    VAR state: StateArray;
        round: CARDINAL;

    BEGIN
        state := in;
        AddRoundKey(state, w[Nr]);          (* See Sec. 5.1.4   *)
        FOR round := Nr-1 TO 1 BY -1 DO
            InvShiftRows(state);            (* See Sec. 5.3.1   *)
            InvSubBytes(state);             (* See Sec. 5.3.2   *)
            AddRoundKey(state, w[round]);
            InvMixColumns(state);           (* See Sec. 5.3.3   *)
        END (*FOR*);
        InvShiftRows(state);
        InvSubBytes(state);
        AddRoundKey(state, w[0]);
        out := state;
    END InvCipher;

(************************************************************************)
(*                            KEY EXPANSION                             *)
(************************************************************************)

PROCEDURE KeyExpansion (Nk: CARDINAL;  VAR (*IN*) key: ARRAY OF CARD8;
                                        VAR (*OUT*) w: ARRAY OF Word);

    (* The input key is 4*Nk bytes long (Nk words). *)
    (* The output array w has Nb*(Nr+1) words.      *)
    (* The first Nk words are just a copy of w.     *)

    (************************************************************)
    (*              Nk      Nr      length of w (words)         *)
    (*  AES-128     4       10              44                  *)
    (*  AES-192     6       12              52                  *)
    (*  AES-256     8       14              60                  *)
    (************************************************************)

    VAR temp, v1: Word;
        i, Nr: CARDINAL;

    BEGIN
        (* The first Nk words of w are a copy of the key. *)

        Copy (ADR(key), ADR(w), 4*Nk);

        (* The remaining words are combinations of the preceding ones. *)

        Nr := Nk+6;
        i := Nk;
        WHILE i < Nb * (Nr+1) DO
            temp := w[i-1];
            IF i MOD Nk = 0 THEN
                v1.byte := SubWord(RotWord(temp.byte));
                temp.word := IXOR (v1.word, Rcon[i DIV Nk]);
            ELSIF (Nk > 6) AND (i MOD Nk = 4) THEN
                temp.byte := SubWord(temp.byte);
            END (*IF*);
            w[i].word := IXOR (w[i-Nk].word, temp.word);
            INC (i);
        END (*WHILE*);

    END KeyExpansion;

(************************************************************************)
(*             CONVERSION BETWEEN BYTE STREAM AND STATE ARRAY           *)
(************************************************************************)

PROCEDURE Pack (amount: CARDINAL;  VAR (*IN*) input: ARRAY OF CARD8;
                 ipos, rpos: CARDINAL;  VAR (*OUT*) result: StateArray);

    (* Puts amount bytes, starting at input[pos], into result starting  *)
    (* at byte rpos of result.  This allows for the case where result   *)
    (* already holds some unprocessed data.                             *)

    BEGIN
        Copy (ADR(input[ipos]), AddOffset(ADR(result), rpos), amount);
    END Pack;

(************************************************************************)

PROCEDURE Unpack (state: StateArray;  VAR (*OUT*) result: ARRAY OF CARD8;
                                        pos: CARDINAL);

    (* Puts 16 bytes into result, starting at result[pos]. *)

    BEGIN
        Copy (ADR(state), ADR(result[pos]), 16);
    END Unpack;

(************************************************************************)

PROCEDURE IncrementCounter (VAR (*INOUT*) count: Word);

    VAR k: CARDINAL;

    BEGIN
        k := 3;
        LOOP
            IF count.byte[k] < MAX(CARD8) THEN
                INC (count.byte[k]);
                RETURN;
            ELSE
                count.byte[k] := 0;
                IF k = 0 THEN
                    RETURN;
                END (*IF*);
                DEC (k);
            END (*IF*);
        END (*LOOP*);
    END IncrementCounter;

(************************************************************************)

PROCEDURE IncrementFullCounter (VAR (*INOUT*) count: StateArray);

    (* Increments the full 16-byte counter. *)

    VAR j, k: CARDINAL;

    BEGIN
        j := 3;  k := 3;
        LOOP
            IF count[k].byte[j] < MAX(CARD8) THEN
                INC (count[k].byte[j]);
                RETURN;
            ELSE
                count[k].byte[j] := 0;
                IF j > 0 THEN DEC(j);
                ELSIF k > 0 THEN
                    DEC (k);  j := 3;
                ELSE
                    j := 3;  k := 3;
                END (*IF*);
            END (*IF*);
        END (*LOOP*);
    END IncrementFullCounter;

(************************************************************************)
(*              GENERIC ENCRYPTION/DECRYPTION PROCEDURE                 *)
(************************************************************************)

PROCEDURE AESgeneric (ctx: AEScontext;  inputlength: CARDINAL;
                                     VAR (*IN*) input: ARRAY OF CARD8;
                                     VAR (*OUT*) output: ARRAY OF CARD8;
                                     VAR (*INOUT*) outpos: CARDINAL);

    (* The input key is 4*Nk bytes long (Nk words).         *)

    (********************************************************)
    (*              Nk      Nr      length of w (words)     *)
    (*  AES-128     4       10              44              *)
    (*  AES-192     6       12              52              *)
    (*  AES-256     8       14              60              *)
    (********************************************************)

    VAR previnput: StateArray;
        amount, inpos, togo: CARDINAL;

    BEGIN
        inpos := 0;  togo := inputlength;
        WHILE togo > 0 DO

            IF ctx^.blockwaiting THEN

                (* Put out a deferred block. *)

                Unpack (ctx^.state, output, outpos);
                INC (outpos, 16);

            END (*IF*);

            previnput := ctx^.inbuff;

            (* Pack next 16 bytes, if possible, into ctx^.inbuff. *)

            amount := 16 - ctx^.incount;
            IF togo < amount THEN
                amount := togo;
            END (*IF*);
            Pack (amount, input, inpos, ctx^.incount, ctx^.inbuff);
            INC (inpos, amount);
            DEC (togo, amount);
            INC (ctx^.incount, amount);

            IF ctx^.incount = 16 THEN

                (* Perform the encryption. *)

                IF ctx^.decrypt THEN
                    InvCipher (ctx^.Nk+6, ctx^.inbuff, ctx^.state, ctx^.roundkeys);
                    IF ctx^.CBC THEN
                        BlockXOR (previnput, ctx^.state);
                    END (*IF*);
                    ctx^.blockwaiting := ctx^.fullpadding;

                ELSE
                    IF ctx^.CBC THEN
                        BlockXOR (ctx^.state, ctx^.inbuff);
                    END (*IF*);
                    IF ctx^.CTR THEN
                        Cipher (ctx^.Nk+6, ctx^.counter, ctx^.state, ctx^.roundkeys);
                        BlockXOR (ctx^.inbuff, ctx^.state);
                        IncrementCounter (ctx^.counter[3]);
                    ELSIF ctx^.SDCTR THEN
                        Cipher (ctx^.Nk+6, ctx^.counter, ctx^.state, ctx^.roundkeys);
                        BlockXOR (ctx^.inbuff, ctx^.state);
                        IncrementFullCounter (ctx^.counter);
                    ELSE
                        Cipher (ctx^.Nk+6, ctx^.inbuff, ctx^.state, ctx^.roundkeys);
                    END (*IF*);
                END (*IF*);
                ctx^.incount := 0;

                IF NOT ctx^.blockwaiting THEN

                    (* Convert result to an array of bytes. *)

                    Unpack (ctx^.state, output, outpos);
                    INC (outpos, 16);
                END (*IF*);

            END (*IF*);

        END (*WHILE*);

    END AESgeneric;

(************************************************************************)
(*                      THE END-USER PROCEDURES                         *)
(************************************************************************)

PROCEDURE AESopen (decrypt, fullpadding: BOOLEAN;  variant: CARDINAL;
                                        key: ARRAY OF CARD8): AEScontext;

    (* Begins one ciphering operation.  The returned value can then be  *)
    (* used in encryption and decryption operations using this key.     *)
    (* We recognise three versions of the algorithm:                    *)
    (*        variant     key length (bytes)      key length (bits)     *)
    (*          128             16                      128             *)
    (*          192             24                      192             *)
    (*          256             32                      256             *)
    (*                                                                  *)
    (* The two Boolean parameters are:                                  *)
    (*                                                                  *)
    (*      decrypt     FALSE for encryption, TRUE for decryption       *)
    (*                                                                  *)
    (*      fullpadding TRUE if you want the padding style specified    *)
    (*                  in PKCS #7 (a repeated extra byte count).       *)
    (*                  If you set this to FALSE then the only padding  *)
    (*                  is Nul bytes added to a final partial block,    *)
    (*                  and there is no padding at all if the input     *)
    (*                  length is a multiple of 16.                     *)
    (*                                                                  *)
    (* During decryption, the padding is stripped if the fullpadding    *)
    (* option is specified.  This does mean that the output is one      *)
    (* block behind the input in that case, because we need to know     *)
    (* which block is the last block before removing padding.           *)

    VAR ctx: AEScontext;
        Nk, count, j: CARDINAL;
        w: ARRAY [0..59] OF Word;

    BEGIN
        IF (variant = 128) OR (variant = 192) OR (variant = 256) THEN
            NEW (ctx);
            ctx^.decrypt := decrypt;
            ctx^.CBC := FALSE;
            ctx^.CTR := FALSE;
            ctx^.SDCTR := FALSE;
            ctx^.fullpadding := fullpadding;
            ctx^.blockwaiting := FALSE;
            ctx^.incount := 0;
            Nk := variant DIV 32;
            ctx^.Nk := Nk;
            KeyExpansion (Nk, key, w);

            (* Move the expansion into StateArray format. *)

            FOR count := 0 TO Nk+6 DO
                FOR j := 0 TO 3 DO
                    ctx^.roundkeys[count][j] := w[4*count + j];
                END (*FOR*);
            END (*FOR*);
        ELSE
            ctx := NIL;
        END (*IF*);

        RETURN ctx;

    END AESopen;

(************************************************************************)

PROCEDURE AES_CBCopen (decrypt, fullpadding: BOOLEAN;  variant: CARDINAL;
                                    key, IV: ARRAY OF CARD8): AEScontext;

    (* Like AESopen, but enables the Cipher Block Chaining (CBC) mode   *)
    (* of AES.  The extra parameter IV is the initialisation vector.    *)

    VAR ctx: AEScontext;

    BEGIN
        ctx := AESopen (decrypt, fullpadding, variant, key);
        ctx^.CBC := TRUE;
        IF ctx^.decrypt THEN
            Pack (16, IV, 0, 0, ctx^.inbuff);
        ELSE
            Pack (16, IV, 0, 0, ctx^.state);
        END (*IF*);
        RETURN ctx;
    END AES_CBCopen;

(************************************************************************)

PROCEDURE AES_CTRopen (variant: CARDINAL;  key: ARRAY OF CARD8;
                        nonce: CARDINAL;  IV: ARRAY OF CARD8): AEScontext;

    (* Like AESopen, but enables the Counter (CTR) mode of AES.  The    *)
    (* nonce (32 bits) and IV (64 bits) are assumed to be already in    *)
    (* bigendian order.                                                 *)

    VAR ctx: AEScontext;
        countwords: ARRAY [0..3] OF CARD32;

    BEGIN
        ctx := AESopen (FALSE, FALSE, variant, key);
        ctx^.CTR := TRUE;
        countwords[0] := nonce;
        Copy (ADR(IV), ADR(countwords[1]), 8);
        countwords[3] := 0;
        Copy (ADR(countwords), ADR(ctx^.counter), 16);
        IncrementCounter (ctx^.counter[3]);
        RETURN ctx;
    END AES_CTRopen;

(************************************************************************)

PROCEDURE AES_SDCTRopen (variant: CARDINAL;  key: ARRAY OF CARD8;
                                    IV: ARRAY OF CARD8): AEScontext;

    (* Similar to AES_CTRopen, but in this case the counter is a full   *)
    (* 16 bytes long, and IV, assumed to be in bigendian order, is the  *)
    (* initial counter value.                                           *)

    VAR ctx: AEScontext;

    BEGIN
        ctx := AESopen (FALSE, FALSE, variant, key);
        ctx^.SDCTR := TRUE;
        Copy (ADR(IV), ADR(ctx^.counter), 16);
        RETURN ctx;
    END AES_SDCTRopen;

(************************************************************************)

PROCEDURE AESprocess (ctx: AEScontext;  inputlength: CARDINAL;
                             VAR (*IN*) input: ARRAY OF CARD8;
                             VAR (*OUT*) output: ARRAY OF CARD8;
                             VAR (*INOUT*) outpos: CARDINAL);

    (* Encrypts or decrypts the input array, depending on what was      *)
    (* specified in AESopen.  The result is stored beginning at         *)
    (* output[outpos], and outpos is updated.  It is legal to call this *)
    (* multiple times before the AESfinal call that terminates the      *)
    (* operation.                                                       *)

    BEGIN
        AESgeneric (ctx, inputlength, input, output, outpos);
    END AESprocess;

(************************************************************************)

PROCEDURE AESfinal (VAR (*INOUT*) ctx: AEScontext;
                            VAR (*OUT*) output: ARRAY OF CARD8;
                             VAR (*INOUT*) outpos: CARDINAL);

    (* Processes any remaining unsent data, adds or removes padding as  *)
    (* needed, then terminates the use of the key specified in AESopen. *)

    VAR padbuff: ARRAY [0..15] OF CARD8;
        j, padamount: CARDINAL;  paddata: CARD8;  error: BOOLEAN;

    BEGIN
        IF ctx <> NIL THEN
            IF ctx^.decrypt THEN

                (* DECRYPTION *)

                IF ctx^.blockwaiting THEN

                    (* This block contains padding, so we have to work  *)
                    (* out how much, if any, is data we have to keep.   *)

                    Unpack (ctx^.state, padbuff, 0);
                    paddata := padbuff[15];
                    error := paddata > 16;
                    IF NOT error THEN
                        FOR j := 16 - paddata TO 15 DO
                            error := error OR (padbuff[j] <> paddata);
                        END (*FOR*);
                    END (*IF*);

                    (* If error, put out the whole block.       *)
                    (* Otherwise, send bytes 0..16-paddata-1.   *)

                    IF error THEN
                        j := 16;
                    ELSE
                        j := 16 - paddata;
                    END (*IF*);
                    IF j > 0 THEN
                        Copy (ADR(padbuff), ADR(output[outpos]), j);
                        INC (outpos, j);
                    END (*IF*);

                END (*IF*);
            ELSE

                (* ENCRYPTION *)

                IF ctx^.CTR THEN

                    (* We don't do padding in the CTR case. *)

                    IF ctx^.incount > 0 THEN

                        (* Do a whole block cipher, but put out only    *)
                        (* ctx^.incount bytes of the result.            *)

                        Cipher (ctx^.Nk+6, ctx^.counter, ctx^.state, ctx^.roundkeys);
                        BlockXOR (ctx^.inbuff, ctx^.state);
                        Copy (ADR(ctx^.state), ADR(output[outpos]), ctx^.incount);
                        INC (outpos, ctx^.incount);

                    END (*IF*);

                ELSIF ctx^.SDCTR THEN

                    (* We don't do padding in the SDCTR case. *)

                    IF ctx^.incount > 0 THEN

                        (* Do a whole block cipher, but put out only    *)
                        (* ctx^.incount bytes of the result.            *)

                        Cipher (ctx^.Nk+6, ctx^.counter, ctx^.state, ctx^.roundkeys);
                        BlockXOR (ctx^.inbuff, ctx^.state);
                        Copy (ADR(ctx^.state), ADR(output[outpos]), ctx^.incount);
                        INC (outpos, ctx^.incount);

                    END (*IF*);

                ELSE
                    (* Add the padding and process final block. *)

                    paddata := 0;
                    padamount := 16 - ctx^.incount;
                    IF ctx^.fullpadding THEN paddata := padamount;
                    ELSIF padamount = 16 THEN padamount := 0;
                    END (*IF*);
                    IF padamount > 0 THEN
                        FOR j := 0 TO padamount-1 DO
                            padbuff[j] := paddata;
                        END (*FR*);
                        AESgeneric (ctx, padamount, padbuff, output, outpos);
                    END (*IF*);
                END (*IF*);
            END (*IF*);

            DISPOSE (ctx);
        END (*IF*);
    END AESfinal;

(************************************************************************)

END AES.

