(**************************************************************************)
(*                                                                        *)
(*  TLS module                                                            *)
(*  Copyright (C) 2025   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       *)
(*                                                                        *)
(**************************************************************************)

IMPLEMENTATION MODULE TLSCrypt;

        (********************************************************)
        (*                                                      *)
        (*  Interface between TLS and the encryption algorithms *)
        (*                                                      *)
        (*  Programmer:         P. Moylan                       *)
        (*  Started:            5 April 2018                    *)
        (*  Last edited:        15 October 2025                 *)
        (*  Status:             Apparently working              *)
        (*                                                      *)
        (********************************************************)


FROM SYSTEM IMPORT CARD8, CARD16, ADR, LOC;

IMPORT AES, DES, RC4, Strings;

FROM VarStrings IMPORT
    (* type *)  ByteStringPtr;

FROM TLSBase IMPORT
    (* proc *)  NYI,
    (* proc *)  TextType;

FROM MiscFuncs IMPORT
    (* proc *)  Swap2, Swap4;

FROM RandCard IMPORT
    (* proc *)  Randomize, RandCardinal;

FROM TransLog IMPORT
    (* type *)  TransactionLogID,
    (* proc *)  LogTransactionL;

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

FROM Alloc IMPORT
    (* type *)  ALLOCATE, DEALLOCATE;

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

TYPE
    ciphertype = (streamcipher, blockcipher, AEADcipher);

    (* See RFC 5246, Appendix C, for how cipher suites decompose into   *)
    (* three critical properties: key exchange, bulk cipher, and MAC.   *)
    (* There is obviously extensive use of suites that are not          *)
    (* mentioned in RFC 5246, and for those I am relying on web sources.*)

VAR
    BlockSize, KeySize: ARRAY BulkCipherAlgorithm OF CARDINAL;

(************************************************************************)
(*                        CHANGING ENCRYPTION RULE                      *)
(************************************************************************)

PROCEDURE InitEncryptionRule (VAR (*OUT*) ER: EncrRule);

    (* Creates an initial "no encryption" state. *)

    BEGIN
        ER.bulkcipher := nocipher;
    END InitEncryptionRule;

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

PROCEDURE SetEncryption (VAR (*INOUT*) rule: EncrRule;  cipher: BulkCipherAlgorithm;
                            VAR (*IN*) key: ARRAY OF CARD8;  keylength: CARDINAL);

    (* Switches to a new encryption algorithm.  *)

    BEGIN
        rule.bulkcipher := cipher;
        Copy (ADR(key), ADR(rule.key), keylength);
    END SetEncryption;

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

PROCEDURE CloseEncryptionRule (VAR (*INOUT*) ER: EncrRule);

    (* Discards the EncrRule data. *)

    (* Not sure whether this is needed. *)

    BEGIN
        ER.bulkcipher := nocipher;
    END CloseEncryptionRule;

(************************************************************************)
(*                            CIPHER PROPERTIES                         *)
(************************************************************************)

PROCEDURE AppendCipherName (cipher: BulkCipherAlgorithm;
                                VAR (*INOUT*) message: ARRAY OF CHAR);

    (* Appends the name of the cipher to the message. *)

    VAR name: ARRAY[0..15] OF CHAR;

    BEGIN
        CASE cipher OF
            | nocipher:      name := "nocipher";
            | rc4_128:       name := "rc4_128";
            | tripledes_cbc: name := "tripledes_cbc";
            | aes128_cbc:    name := "aes128_cbc";
            | aes256_cbc:    name := "aes256_cbc";
        ELSE
            name := "unknown";
        END (*IF*);
        Strings.Append (name, message);
    END AppendCipherName;

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

PROCEDURE KeyLength (cipher: BulkCipherAlgorithm): CARDINAL;

    (* Returns the key length in bytes. *)

    BEGIN
        RETURN KeySize[cipher];
    END KeyLength;

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

PROCEDURE IVLength (cipher: BulkCipherAlgorithm): CARDINAL;

    (* Returns the initialisation vector length in bytes.  This is  *)
    (* for block ciphers only.                                      *)

    BEGIN
        CASE cipher OF
            | nocipher, rc4_128:
                RETURN 0;
            | tripledes_cbc, aes128_cbc, aes256_cbc:
                RETURN BlockSize[cipher];
        END (*CASE*);
    END IVLength;

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

PROCEDURE FixedIVLength (cipher: BulkCipherAlgorithm): CARDINAL;

    (* Returns the initialisation vector length in bytes.  This is  *)
    (* for AEAD ciphers only.                                       *)

    BEGIN
        CASE cipher OF
            | nocipher, rc4_128, tripledes_cbc, aes128_cbc, aes256_cbc:
                RETURN 0;
        END (*CASE*);
    END FixedIVLength;

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

PROCEDURE CipherType (ER: EncrRule): ciphertype;

    (* Is ER a stream, block, or AEAD cipher? *)

    BEGIN
        CASE ER.bulkcipher OF
            | nocipher, rc4_128:                        RETURN streamcipher;
            | tripledes_cbc, aes128_cbc, aes256_cbc:    RETURN blockcipher;
        END (*CASE*);
    END CipherType;

(************************************************************************)
(*                          ENCRYPTION/DECRYPTION                       *)
(************************************************************************)

PROCEDURE DoEncryption (decrypt: BOOLEAN;  rule: EncrRule;
                            VAR (*INOUT*) dataptr: ByteStringPtr;
                              offset, size: CARDINAL);

    (* This does the actual encryption or decryption step.  We encrypt  *)
    (* the bytes from dataptr[offset] to dataptr^[offset+size-1].  The  *)
    (* first offset bytes hold the IV, which is not encrypted.  (If     *)
    (* offset = 0, there is no IV.)  The result replaces the original   *)
    (* data, and has the same size as the original.                     *)

    VAR desctx: DES.DEScontext;  aesctx: AES.AEScontext;  rc4ctx: RC4.RC4state;
        cipher: BulkCipherAlgorithm;
        IV: ARRAY [0..MaxBlockSize-1] OF CARD8;
        pos, variant: CARDINAL;
        p, q: ByteStringPtr;

    BEGIN
        IF rule.bulkcipher = nocipher THEN
            RETURN;
        END (*IF*);
        p := dataptr;
        IF offset > 0 THEN
            Copy (p, ADR(IV), offset);
            p := AddOffset (p, offset);
        END (*IF*);
        ALLOCATE (q, size);
        IF q = NIL THEN
            RETURN;
        END (*IF*);

        (* Note that size is the amount of data to be encrypted or  *)
        (* decrypted. It does not include the size of IV.           *)

        cipher := rule.bulkcipher;
        IF cipher = aes128_cbc THEN
            variant := 128;
        ELSE
            variant := 256;
        END (*IF*);
        pos := 0;
        CASE cipher OF
            |   rc4_128:
                    rc4ctx := RC4.RC4Init (0, rule.key, 16);
                    RC4.Encrypt (rc4ctx, size, p^, q^);
                    RC4.RC4Close (rc4ctx);
            |   tripledes_cbc:
                    desctx := DES.DES_CBCopen (decrypt, FALSE, 24, rule.key, IV);
                    DES.DESEncrypt (desctx, size, p^, q^, pos);
                    DES.DESfinal (desctx, q^, pos);
            |   aes128_cbc, aes256_cbc:
                    aesctx := AES.AES_CBCopen (decrypt, FALSE, variant, rule.key, IV);
                    AES.AESprocess (aesctx, size, p^, q^, pos);
                    AES.AESfinal (aesctx, q^, pos);
        END (*CASE*);

        (* Now let the result replace the original data.  (It has the   *)
        (* same size, so we can just overlay it.)  Of course we should  *)
        (* deallocate the temporary variable q^.                        *)

        Copy (q, p, size);
        DEALLOCATE (q, size);     (* used to cause a crash *)

    END DoEncryption;

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

PROCEDURE GenRandom (VAR (*OUT*) IV: ARRAY OF LOC;  length: CARDINAL);

    (* Generates a random string of length bytes. *)

    VAR j: CARDINAL;  IVC: ARRAY [0..7] OF CARDINAL;

    BEGIN
        (* My method here is a little quick-and-dirty, but I can update *)
        (* it later.  This generates length MOD 4 surplus bytes that    *)
        (* are not used.  The caller will then produce at least this    *)
        (* many bytes that will overlay the surplus.                    *)

        FOR j := 0 TO length DIV 4 DO
            IVC[j] := RandCardinal();
        END (*FOR*);
        Copy (ADR(IVC), ADR(IV), length);
    END GenRandom;

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

PROCEDURE Encrypt (VAR (*IN*) rule: EncrRule;  logID: TransactionLogID;
                   VAR (*INOUT*) buffer: TextType;
                   VAR (*IN*) MAC: ARRAY OF CARD8;  maclength: CARDINAL);

    (* We encrypt the message in buffer after appending the MAC and     *)
    (* possibly some padding. The encrypted string replaces the         *)
    (* original message.                                                *)

    (* In all cases we leave the type, version, and length unencrypted, *)
    (* although we might change buffer.length.  For stream ciphers we   *)
    (* encrypt data followed by mac.  For block ciphers we start with   *)
    (* IV (unencrypted), and then we encrypt the content followed by    *)
    (* mac followed by padding followed by padding length.              *)

    (* I haven't implemented any AEAD ciphers.  *)

    VAR ctype: ciphertype;  IVsize, size, blocksize, L, j: CARD16;
        padlength: CARD8;
        p, q: ByteStringPtr;

    BEGIN
        padlength := 0;  IVsize := 0;
        IF (rule.bulkcipher = nocipher) AND (maclength = 0) THEN

            (* Special case: if no encryption and no MAC is appended,   *)
            (* we just return the original buffer contents.             *)

        ELSE
            (* Work out how much space to allocate. *)

            ctype := CipherType (rule);
            size := Swap2 (buffer.length);
            IF ctype = streamcipher THEN INC(size, maclength);
            ELSIF ctype = blockcipher THEN
                blocksize := BlockSize[rule.bulkcipher];
                IVsize := blocksize;
                INC (size, maclength+1);
                padlength := blocksize - (size MOD blocksize);
                INC (size, padlength);
            END (*IF*);

            (* Make a copy of the original data, and then clear buffer  *)
            (* in preparation for putting the encrypted result there.   *)

            ALLOCATE (p, IVsize + size);
            q := p;
            IF IVsize > 0 THEN

               (* Although IV is called an initialisation vector, TLS   *)
               (* is unusual in that a new IV is generated for every packet.*)

                GenRandom (q^, IVsize);
                q := AddOffset (q, IVsize);
            END (*IF*);
            L := Swap2 (buffer.length);
            Copy (buffer.fragment, q, L);
            DEALLOCATE (buffer.fragment, L);
            q := AddOffset (q, L);

            (* Create the rest of the byte string to be encrypted. *)

            IF ctype = streamcipher THEN

                (* Use content followed by mac *)

                IF maclength > 0 THEN
                    Copy (ADR(MAC), q, maclength);
                END (*IF*);

            ELSIF ctype = blockcipher THEN

                (* Content, then mac, then padding *)

                Copy (ADR(MAC), q, maclength);
                q := AddOffset (q, maclength);
                FOR j := 0 TO padlength DO
                    q^[0] := padlength;  q := AddOffset(q, 1);
                END (*FOR*);

            ELSE
                NYI (logID, "AEAD encryption");
            END (*IF*);

            (* Now do the encryption. *)

            DoEncryption (FALSE, rule, p, IVsize, size);

            (* For a stream cipher, p^ is the final answer.  For a      *)
            (* block cipher or an AEAD cipher, we have already put the  *)
            (* IV before the encrypted result.                          *)

            buffer.fragment := p;
            INC (size, IVsize);
            buffer.length := Swap2(size);

        END (*IF*);
    END Encrypt;

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

PROCEDURE Decrypt (VAR (*IN*) rule: EncrRule;  logID: TransactionLogID;
                     VAR (*INOUT*) buffer: TextType;
                      VAR (*OUT*) MAC: ARRAY OF CARD8;  maclength: CARDINAL): BOOLEAN;

    (* Decrypts an N-byte block of data, also returns the mac. *)

    VAR p, q: ByteStringPtr;  L: CARD16;
        blocksize, padlength, j, IVlength: CARDINAL;
        success: BOOLEAN;
        ctype: ciphertype;

    BEGIN
        success := TRUE;

        (* Let p point to the "content" part of the block, and  *)
        (* let L be the length of the content.                  *)

        p := buffer.fragment;  L := Swap2(buffer.length);  IVlength := 0;
        blocksize := BlockSize[rule.bulkcipher];
        ctype := CipherType (rule);
        IF ctype = streamcipher THEN
            IF (rule.bulkcipher = nocipher) AND (maclength = 0) THEN
                (* Special case -- data not encrypted, no MAC. *)
                RETURN TRUE;
            END (*IF*);
        ELSE
            (* The block starts with an unencrypted IV.  *)

            IVlength := blocksize;
            DEC (L, IVlength);
        END (*IF*);

        (* Now the decryption step. *)

        DoEncryption (TRUE, rule, p, IVlength, L);
        IF IVlength > 0 THEN
            p := AddOffset (p, IVlength);
        END (*IF*);

        IF ctype = AEADcipher THEN
            NYI (logID, "AEAD decryption");
        ELSIF ctype = blockcipher THEN

            (* Remove the padding. *)

            padlength := p^[L-1];  j := padlength+1;
            REPEAT
                DEC (L);  DEC(j);
                IF p^[L] <> padlength THEN
                    LogTransactionL (logID, "ERROR: incorrect padding value");
                    success := FALSE;

                    (* If the padding value is wrong, there is no way   *)
                    (* to know what L should be, so we should just      *)
                    (* leave it alone and declare an error.             *)

                    j := 0;
                END (*IF*);
            UNTIL j = 0;

        END (*IF*);

        (* Extract the MAC. *)

        IF maclength > 0 THEN
            Copy (ADR(p^[L-maclength]), ADR(MAC), maclength);
            DEC (L, maclength);
        END (*IF*);

        (* Discard the original packet, and replace it with the decrypted result. *)

        IF L > 0 THEN
            ALLOCATE (q, L);
            IF q = NIL THEN
                LogTransactionL (logID, "Out of heap memory");
                success := FALSE;
            ELSE
                Copy (p, q, L);
            END (*IF*);
        ELSE
            q := NIL;
        END (*IF*);
        DEALLOCATE (buffer.fragment, Swap2(buffer.length));
        buffer.fragment := q;
        buffer.length := Swap2(L);
        RETURN success;

    END Decrypt;

(************************************************************************)
(*             INITIALISING THE BlockSize AND KeySize ARRAYS            *)
(************************************************************************)

PROCEDURE InitSizes;

    BEGIN
        BlockSize[nocipher] := 1;
        BlockSize[rc4_128] := 1;
        BlockSize[tripledes_cbc] := 8;
        BlockSize[aes128_cbc] := 16;
        BlockSize[aes256_cbc] := 16;

        KeySize[nocipher] := 0;
        KeySize[rc4_128] := 16;
        KeySize[tripledes_cbc] := 24;
        KeySize[aes128_cbc] := 16;
        KeySize[aes256_cbc] := 32;

    END InitSizes;

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

BEGIN
    InitSizes;
    Randomize;
END TLSCrypt.

