sm4/aes_FPE ，FF3-1 保留格式加密

"""
Author: tanglei
DateTime:2024-11-18 完成
微信：ciss_cedar
欢迎一起学习
使用了其他人的方法封装
"""

# Package ff3 implements the FF3-1 format-preserving encryption algorithm/scheme

import math
import string

from symmetric_alg import MySymCipher, AlgName, CipherMode

# The recommendation in Draft SP 800-38G was strengthened to a requirement in Draft
# SP 800-38G Revision 1: the minimum domain size for FF1 and FF3-1 is one million.

NUM_ROUNDS = 8
BLOCK_SIZE = 16  # aes.BlockSize
TWEAK_LEN = 8  # Original FF3 tweak length
TWEAK_LEN_NEW = 7  # FF3-1 tweak length
HALF_TWEAK_LEN = TWEAK_LEN // 2


def reverse_string(txt):
    """func defined for clarity"""
    return txt[::-1]


class FF3Cipher:

    DOMAIN_MIN = 1_000_000  # 1M required in FF3-1
    BASE_STR = string.digits+ string.ascii_uppercase + string.ascii_lowercase
    Other='+/='
    BASE_STR=BASE_STR+Other
    BASE_STR_LEN = len(BASE_STR)
    RADIX_MAX = 256  # Support 8-bit alphabets for now
    # radix=10 数字， 16 16进制字符，65 base64 = 用于填充故实际上是65个字符
    def __init__(self, key,iv,mode,alg_name,tweak, radix=10):
        key_bytes = bytes.fromhex(key)
        self.key=key
        self.iv=iv
        self.mode=mode
        self.alg_name=alg_name
        self.tweak = str(tweak).ljust(16,'0')
        self.radix = radix
        if radix <= FF3Cipher.BASE_STR_LEN:
            self.alphabet = FF3Cipher.BASE_STR[0:radix]
        else:
            self.alphabet = None

        # Calculate range of supported message lengths [minLen..maxLen]
        # per revised spec, radix^minLength >= 1,000,000.
        self.minLen = math.ceil(math.log(FF3Cipher.DOMAIN_MIN) / math.log(radix))

        # We simplify the specs log[radix](2^96) to 96/log2(radix) using the log base
        # change rule
        self.maxLen = 2 * math.floor(96/math.log2(radix))

        klen = len(key_bytes)

        # Check if the key is 128, 192, or 256 bits = 16, 24, or 32 bytes
        if klen not in (16, 24, 32):
            raise ValueError(f'key length is {klen} but must be 128, 192, or 256 bits')

        # While FF3 allows radices in [2, 2^16], commonly useful range is 2..62
        if (radix < 2) or (radix > FF3Cipher.RADIX_MAX):
            raise ValueError("radix must be between 2 and 62, inclusive")

        # Make sure 2 <= minLength <= maxLength
        if (self.minLen < 2) or (self.maxLen < self.minLen):
            raise ValueError("minLen or maxLen invalid, adjust your radix")

        # AES block cipher in ECB mode with the block size derived based on the length
        # of the key. Always use the reversed key since Encrypt and Decrypt call ciph
        # expecting that

        #self.aesCipher = AES.new(reverse_string(key_bytes), AES.MODE_ECB)

    # factory method to create a FF3Cipher object with a custom alphabet
    @staticmethod
    def withCustomAlphabet(key, tweak, alphabet):
        c = FF3Cipher(key, tweak, len(alphabet))
        c.alphabet = alphabet
        return c

    def encrypt(self, plaintext):
        """Encrypts the plaintext string and returns a ciphertext of the same length
        and format"""
        return self.encrypt_with_tweak(plaintext, self.tweak)


    # EncryptWithTweak allows a parameter tweak instead of the current Cipher's tweak

    def encrypt_with_tweak(self, plaintext, tweak):

        tweakBytes = bytes.fromhex(tweak)

        n = len(plaintext)

        # Check if message length is within minLength and maxLength bounds
        if (n < self.minLen) or (n > self.maxLen):
            raise ValueError(f"message length {n} is not within min {self.minLen} and "
                             f"max {self.maxLen} bounds")

        # Make sure the given the length of tweak in bits is 56 or 64
        if len(tweakBytes) not in [TWEAK_LEN, TWEAK_LEN_NEW]:
            raise ValueError(f"tweak length {len(tweakBytes)} invalid: tweak must be 56"
                             f" or 64 bits")

        # Todo: Check message is in current radix

        # Calculate split point
        u = math.ceil(n / 2)
        v = n - u

        # Split the message
        A = plaintext[:u]
        B = plaintext[u:]

        if len(tweakBytes) == TWEAK_LEN_NEW:
            # FF3-1
            tweakBytes = calculate_tweak64_ff3_1(tweakBytes)

        Tl = tweakBytes[:HALF_TWEAK_LEN]
        Tr = tweakBytes[HALF_TWEAK_LEN:]
       # logger.debug(f"Tweak: {tweak}, tweakBytes:{tweakBytes.hex()}")

        # Pre-calculate the modulus since it's only one of 2 values,
        # depending on whether i is even or odd

        modU = self.radix ** u
        modV = self.radix ** v
       # logger.debug(f"modU: {modU} modV: {modV}")

        # Main Feistel Round, 8 times
        #
        # AES ECB requires the number of bits in the plaintext to be a multiple of
        # the block size. Thus, we pad the input to 16 bytes

        for i in range(NUM_ROUNDS):
            # logger.debug(f"-------- Round {i}")
            # Determine alternating Feistel round side
            if i % 2 == 0:
                m = u
                W = Tr
            else:
                m = v
                W = Tl

            # P is fixed-length 16 bytes
            P = calculate_p(i, self.alphabet, W, B)
            revP = reverse_string(P)

            key = bytes.fromhex(self.key)
            iv =  bytes.fromhex(self.iv)
            my_alg = MySymCipher(key, iv, self.mode, self.alg_name)
            S = my_alg.encrypt_bytes(bytes(revP))

           # S = self.aesCipher.encrypt(bytes(revP))

            S = reverse_string(S)
            # logger.debug("S:    ", S.hex())

            y = int.from_bytes(S, byteorder='big')

            # Calculate c
            c = decode_int_r(A,  self.alphabet)

            c = c + y

            if i % 2 == 0:
                c = c % modU
            else:
                c = c % modV

            # logger.debug(f"m: {m} A: {A} c: {c} y: {y}")
            C = encode_int_r(c, self.alphabet, int(m))

            # Final steps
            A = B
            B = C

            # logger.debug(f"A: {A} B: {B}")

        return A + B

    def decrypt(self, ciphertext):

        return self.decrypt_with_tweak(ciphertext, self.tweak)

    def decrypt_with_tweak(self, ciphertext, tweak):

        tweakBytes = bytes.fromhex(tweak)

        n = len(ciphertext)

        # Check if message length is within minLength and maxLength bounds
        if (n < self.minLen) or (n > self.maxLen):
            raise ValueError(f"message length {n} is not within min {self.minLen} and "
                             f"max {self.maxLen} bounds")

        # Make sure the given the length of tweak in bits is 56 or 64
        if len(tweakBytes) not in [TWEAK_LEN, TWEAK_LEN_NEW]:
            raise ValueError(f"tweak length {len(tweakBytes)} invalid: tweak must be 8 "
                             f"bytes, or 64 bits")

        # Todo: Check message is in current radix

        # Calculate split point
        u = math.ceil(n/2)
        v = n - u

        # Split the message
        A = ciphertext[:u]
        B = ciphertext[u:]

        if len(tweakBytes) == TWEAK_LEN_NEW:
            # FF3-1
            tweakBytes = calculate_tweak64_ff3_1(tweakBytes)

        Tl = tweakBytes[:HALF_TWEAK_LEN]
        Tr = tweakBytes[HALF_TWEAK_LEN:]
        #logger.debug(f"Tweak: {tweak}, tweakBytes:{tweakBytes.hex()}")

        # Pre-calculate the modulus since it's only one of 2 values,
        # depending on whether i is even or odd

        modU = self.radix ** u
        modV = self.radix ** v
        #logger.debug(f"modU: {modU} modV: {modV}")

        # Main Feistel Round, 8 times

        for i in reversed(range(NUM_ROUNDS)):

            # logger.debug(f"-------- Round {i}")
            # Determine alternating Feistel round side
            if i % 2 == 0:
                m = u
                W = Tr
            else:
                m = v
                W = Tl

            # P is fixed-length 16 bytes
            P = calculate_p(i, self.alphabet, W, A)
            revP = reverse_string(P)

           # S = self.aesCipher.encrypt(bytes(revP))
            key = bytes.fromhex(self.key)
            #alg_name = AlgName.SM4.value
            #mode = CipherMode.CBC.value
            iv = bytes.fromhex(self.iv)
            my_alg = MySymCipher(key, iv, self.mode, self.alg_name)
            S = my_alg.encrypt_bytes(bytes(revP))
            #plain = my_alg.decrypt_bytes(cipher)
            S = reverse_string(S)

            # logger.debug("S:    ", S.hex())

            y = int.from_bytes(S, byteorder='big')

            # Calculate c
            c = decode_int_r(B, self.alphabet)

            c = c - y

            if i % 2 == 0:
                c = c % modU
            else:
                c = c % modV

            # logger.debug(f"m: {m} B: {B} c: {c} y: {y}")
            C = encode_int_r(c, self.alphabet, int(m))

            # Final steps
            B = A
            A = C

            # logger.debug(f"A: {A} B: {B}")

        return A + B


def calculate_p(i, alphabet, W, B):
    # P is always 16 bytes
    P = bytearray(BLOCK_SIZE)

    P[0] = W[0]
    P[1] = W[1]
    P[2] = W[2]
    P[3] = W[3] ^ int(i)

    # The remaining 12 bytes of P are for rev(B) with padding

    BBytes = decode_int_r(B, alphabet).to_bytes(12, "big")
    # logger.debug(f"B: {B} BBytes: {BBytes.hex()}")

    P[BLOCK_SIZE - len(BBytes):] = BBytes
    return P

def calculate_tweak64_ff3_1(tweak56):
    tweak64 = bytearray(8)
    tweak64[0] = tweak56[0]
    tweak64[1] = tweak56[1]
    tweak64[2] = tweak56[2]
    tweak64[3] = (tweak56[3] & 0xF0)
    tweak64[4] = tweak56[4]
    tweak64[5] = tweak56[5]
    tweak64[6] = tweak56[6]
    tweak64[7] = ((tweak56[3] & 0x0F) << 4)
    return tweak64

def encode_int_r(n, alphabet, length=0):

    base = len(alphabet)
    if (base > FF3Cipher.RADIX_MAX):
        raise ValueError(f"Base {base} is outside range of supported radix "
                         f"2..{FF3Cipher.RADIX_MAX}")

    x = ''
    while n >= base:
        n, b = divmod(n, base)
        x += alphabet[b]
    x += alphabet[n]

    if len(x) < length:
        x = x.ljust(length, alphabet[0])

    return x


def decode_int_r(astring, alphabet):

    strlen = len(astring)
    base = len(alphabet)
    num = 0

    idx = 0
    try:
        for char in reversed(astring):
            power = (strlen - (idx + 1))
            num += alphabet.index(char) * (base ** power)
            idx += 1
    except ValueError:
        raise ValueError(f'char {char} not found in alphabet {alphabet}')

    return num

# def encrypt_idcard(key, iv, mode, alg_name, tweak, idcard,radix=10):
#     #alg_name = AlgName.SM4.value
#     #mode = CipherMode.CBC.value
#     temp_source=''
#     if radix == 10 and len(idcard) == 18:
#         if idcard[-1:] in ('X','x'):
#             temp_source = idcard[:17]+'2'
#         else:
#             temp_source=idcard
#         myFF3Cipher = FF3Cipher(key, iv, mode, alg_name, tweak, radix=radix)
#         result_value = myFF3Cipher.encrypt(temp_source)
#         return result_value
#     else:
#             return 'not id_card'
#
# def decrypt_idcard(key, iv, mode, alg_name, tweak, enc_idcard,radix=10,):
#         if radix == 10 and len(enc_idcard) == 18:
#             myFF3Cipher = FF3Cipher(key, iv, mode, alg_name, tweak, radix=radix)
#             result_value = myFF3Cipher.decrypt(enc_idcard)
#             if result_value[-1:]=='2':
#                 result_value=result_value[0:17]+'X'
#             return result_value
#         else:
#             return 'not id_card'

def main():
    # 示例用法


    key = '2934412A66B7A186DC35DC40E926F9EE'
    iv = '86CD720D75F4622DBE96078A3CD1076E'
    # key = bytes.fromhex(key)
    # iv = bytes.fromhex(iv)
    tweak = '86CD720D75F4622D'

    radix = 65  # 基数
    plaintext = '12345678X/'

    alg_name = AlgName.SM4.value
    mode = CipherMode.CBC.value
    myFF3Cipher = FF3Cipher(key, iv, mode, alg_name, tweak, radix=radix)
    ciphertext = myFF3Cipher.encrypt(plaintext)
    print("plaintext:", plaintext)
    print("Ciphertext:", ciphertext)
    decrypted = myFF3Cipher.decrypt(ciphertext)
    print("Decrypted:", decrypted)
    print('-'*66)
    id_card = '10110219791119002X'
    radix=10
    if radix == 10 and len(id_card) == 18:
        if id_card[-1:] in ('X','x'):
                plaintext = id_card[:17]+'2'
    myFF3Cipher = FF3Cipher(key, iv, mode, alg_name, tweak, radix=radix)
    enc_id_card = myFF3Cipher.encrypt(plaintext)

    dec_id_card=myFF3Cipher.decrypt(enc_id_card)
    if dec_id_card[-1:]=='2':
        dec_id_card=dec_id_card[0:17]+'X'
    print(f'id_card={id_card}')
    print(f'enc_id_card={enc_id_card}')
    print(f'dec_id_card={dec_id_card}')

if __name__ == "__main__":
    main()