from Crypto.Cipher import AES
from Crypto.Util import Counter  

# La fonction "printhex" affiche un tableau de bytes en format hexadécimal.
def printhex(x, sep = ' '):
    # On initialise une chaîne vide nommée "str".
    str = ''
    # Pour chaque octet (b) dans x.
    for b in x:  
        # On convertit l'octet en hexadécimal (sans le préfixe '0x').
        byte = hex(b)[2:]
        # Si la longueur de l'octet converti en hexadécimal est inférieure à 2.
        if (len(byte)<2):
            # On ajoute '0' devant l'octet pour le compléter à 2 caractères, suivi du séparateur.
            str += '0' + byte + sep
        else:
            # Sinon, on ajoute l'octet converti en hexadécimal suivi du séparateur.
            str += hex(b)[2:] + sep 
    # On affiche la chaîne, en supprimant le dernier séparateur et en mettant tout en majuscules.
    print (str[:-1].upper())

# La fonction "returnhex" retourne un tableau de bytes en format hexadécimal.
def returnhex(x, sep = ' '):
    str = ''
    for b in x:
        byte = hex(b)[2:]
        if (len(byte)<2):
            str += '0' + byte + sep
        else:
            str += hex(b)[2:] + sep 
    # On retourne la chaîne, en supprimant le dernier séparateur et en mettant tout en majuscules.
    return (str[:-1].upper())

# La fonction "pad" ajoute des octets 0x00 à un tableau de bytes jusqu'à ce qu'il soit de longueur multiple de 16.
def pad(x):
    # On calcule le nombre d'octets nécessaires pour compléter x à une longueur multiple de 16.
    n=(16-len(x)%16)%16
    # On retourne le tableau original x avec les octets 0x00 ajoutés.
    return x + bytes([0x00]*n)


# Fonction de déchiffrement d'une trame
def decrypt(trame):  
  print(trame) # Affiche la trame
  
  # Clé utilisée pour le chiffrement
  key = bytes([0xAD, 0x8E, 0xBB, 0xC4, 0xF9, 0x6A, 0xE7, 0x00, 0x05, 0x06, 0xD3, 0xFC, 0xD1, 0x62, 0x7F, 0xB8])

  # Variables L et M utilisées plus tard
  L = 2
  M = 4

  # Récupération de l'en-tête réseau (NwkHeader) à partir de la trame reçue
  NwkHeader = bytes.fromhex(trame[18:34])

  # Récupération de l'en-tête auxiliaire (AuxHeader) à partir de la trame reçue
  AuxiliaryHeader = bytes.fromhex('2D' + trame[36:62])

  # Formation du nonce à partir de différentes parties de la trame reçue
  nonce = bytes.fromhex(trame[44:60]+trame[36:44]+'2D')

  # Chaîne d'octets "a"
  a = NwkHeader + AuxiliaryHeader

  # Chaîne d'octets "m"
  m = bytes.fromhex(trame[62:94])

  # Concaténation de la longueur de "a" avec la chaîne d'octets "a" elle-même
  AddAuthData = len(a).to_bytes(2, byteorder = 'big') + a
  # Formation de la chaîne d'octets "AddAuthData" remplie
  AddAuthData = pad(AddAuthData)

  # Chaîne d'octets "PlaintextData"
  PlaintextData = m
  # Remplissage (pas nécessaire ici, car la longueur de "m" = 16)
  PlaintextData = pad(PlaintextData)

  # Concaténation de "AddAuthData" et "PlaintextData" pour former la chaîne d'octets "AuthData"
  AuthData = AddAuthData + PlaintextData

  # Définition des drapeaux
  Flags = bytes([0x49]) # = ([0b01001001])

  # Formation de la chaîne d'octets "B0"
  B0 = Flags + nonce + len(m).to_bytes(2, byteorder = 'big')

  # Chaîne d'octets "X0" initiale
  X0 = bytes([0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])

  # Chiffrement de "B0 + AuthData" avec "key"
  cipher = AES.new(key, AES.MODE_CBC, X0)
  X1 = cipher.encrypt(B0 + AuthData)
  T = X1[-16:-12]

  # Définition du flag
  Flags = bytes ([0b00000001])

  # Construction de A0 et A1
  A0 = Flags + nonce + bytes([0x00, 0x00])
  A1 = Flags + nonce + bytes([0x00, 0x01])

  # Chiffrement du PlaintextData avec le mode CTR
  from Crypto.Util import Counter  
  ctr = Counter.new(128, initial_value = int.from_bytes(A1, "big"))
  cipher = AES.new(key, AES.MODE_CTR, counter = ctr)
  Ciphertext = cipher.encrypt(PlaintextData)

  # Chiffrement de A0 avec le mode ECB
  cipher = AES.new(key, AES.MODE_ECB)
  S0 = cipher.encrypt(A0)

  # Effectue S0[0:4] XOR T
  U = bytes(a ^ b for (a, b) in zip(S0[0:4], T))

  # Retourne Ciphertext en format hexadécimal
  return returnhex(Ciphertext)


def decoupeTrame(trame):
    a = ""
    for i in range(len(trame)):
        if (i%2) == 0 and i>1:
            a = a + " " + trame[i]
        else:
            a = a + "" +trame[i]
    list_trame = a.split()
    for j in range(len(list_trame)):
        if j < len(list_trame):
            list_trame[j] += " "
    return a , list_trame

#print(T_trame,"\n",list_trame)

def incrementTrame(trame):
  T_trame,list_trame = decoupeTrame(trame)
  tab = [2,16,18,38,40]
  for j in range(len(tab)):
    for i in range(len(list_trame)):
      if (i == tab[j]):
        list_trame[i] = int(list_trame[i],16)
        list_trame[i]+= 0x01
        list_trame[i] = str(hex(list_trame[i]))
        list_trame[i] = list_trame[i].replace("0x","")
        list_trame[i] = list_trame[i].upper()
      else:
        list_trame[i] = list_trame[i].replace(" ","")
  return list_trame

def concatenageTrame(trame):
  tramestr = ""
  for i in range(len(trame)):
    tramestr += trame[i]
  return tramestr