Update cryptography library and improve security

extensions/confessions.py

@@ -9,6 +9,7 @@
 from __future__ import annotations
 
 import time
+from base64 import b64encode
 from typing import Optional, Union, TYPE_CHECKING
 import discord
 from discord import app_commands
@@ -50,7 +51,7 @@ def __init__(self, bot:MerelyBot):
     if 'report_channel' not in self.config:
       self.config['report_channel'] = ''
     if 'secret' not in self.config or self.config['secret'] == '':
-      self.config['secret'] = self.crypto.keygen(32)
+      self.config['secret'] = b64encode(self.crypto.srandom_token(32)).decode('ascii')
       if not bot.quiet:
         print(
           " - WARN: Your security key has been regenerated. Old confessions are now incompatible"

extensions/confessions_common.py

@@ -5,9 +5,10 @@
 """
 from __future__ import annotations
 
-import io, os, base64, hashlib, re
-from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
-from cryptography.hazmat.backends import default_backend
+import io, re, secrets
+from base64 import b64encode, b64decode
+from Crypto.Cipher import AES
+from Crypto.Hash import SHA
 from typing import Optional, Union, TYPE_CHECKING
 from collections import OrderedDict
 import discord
@@ -409,6 +410,7 @@ async def on_timeout(self):
 class Crypto:
   """ Handles encryption and decryption of sensitive data """
   _key = None
+  NONCE_LEN = 16 # 128 bits
 
   @property
   def key(self) -> str:
@@ -418,30 +420,37 @@ def key(self) -> str:
   @key.setter
   def key(self, value:str):
     """ Key setter """
-    self._key = base64.decodebytes(bytes(value, encoding='ascii'))
+    self._key = b64decode(value)
 
-  def setup(self, nonce:bytes = b'\xae[Et\xcd\n\x01\xf4\x95\x9c|No\x03\x81\x98'):
+  def setup(self, nonce:bytes):
     """ Initializes the AES-256 scheme """
-    backend = default_backend()
-    cipher = Cipher(algorithms.AES(self._key), modes.CTR(nonce), backend=backend)
-    return (cipher.encryptor(), cipher.decryptor())
+    return AES.new(self._key, AES.MODE_OFB, nonce)
 
-  def keygen(self, length:int = 32) -> str:
-    """ Generates a key for storage """
-    return base64.encodebytes(os.urandom(length)).decode('ascii')
+  def srandom_token(self, length:int = 16) -> bytes:
+    """ Generates a secure random token """
+    return secrets.token_bytes(length)
+
+  def hash(self, data:bytes, salt:bytes) -> bytes:
+    hash = SHA.new(data + salt)
+    return hash.digest()
 
   def encrypt(self, data:bytes) -> bytes:
-    """ Encodes data and returns a base64 string ready for storage """
-    encryptor, _ = self.setup()
-    encrypted = encryptor.update(data) + encryptor.finalize()
-    encodedata = base64.b64encode(encrypted)
-    return encodedata
+    """
+      Encodes data and returns secure bytes for storage
+      The encrypted data will be 16 bytes longer
+    """
+    nonce = self.srandom_token(self.NONCE_LEN)
+    cypher = self.setup(nonce)
+    encrypted = cypher.encrypt(data)
+    encodedata = encrypted
+    return nonce + encodedata
 
-  def decrypt(self, data:str) -> bytes:
+  def decrypt(self, data:bytes) -> bytes:
     """ Read encoded data and return the raw bytes that created it """
-    _, decryptor = self.setup()
-    encrypted = base64.b64decode(data)
-    rawdata = decryptor.update(encrypted) + decryptor.finalize()
+    nonce = data[:self.NONCE_LEN]
+    cypher = self.setup(nonce)
+    encrypted = data[self.NONCE_LEN:]
+    rawdata = cypher.decrypt(encrypted)
     return rawdata
 
 
@@ -451,6 +460,7 @@ def decrypt(self, data:str) -> bytes:
 class ConfessionData:
   """ Dataclass for Confessions """
   SCOPE = 'confessions' # exists to keep babel happy
+  DATA_VERSION = 2
   anonid:str | None
   author:discord.User
   targetchannel:discord.TextChannel
@@ -475,18 +485,20 @@ def __init__(self, parent:Union[Confessions, ConfessionsModeration]):
 
   async def from_binary(self, crypto:Crypto, rawdata:str):
     """ Creates ConfessionData from an encrypted binary string """
-    binary = crypto.decrypt(rawdata)
-    if len(binary) == 24: # TODO: Legacy format, remove eventually
-      author_id = int.from_bytes(binary[0:8], 'big')
-      targetchannel_id = int.from_bytes(binary[16:24], 'big')
-    elif len(binary) == 25:
-      author_id = int.from_bytes(binary[0:8], 'big')
-      targetchannel_id = int.from_bytes(binary[8:16], 'big')
+    binary = crypto.decrypt(b64decode(rawdata))
+    if len(binary) == 26:
+      data_version = int.from_bytes(bytes((binary[0],)), 'big')
+      if data_version != self.DATA_VERSION:
+        raise CorruptConfessionDataException(
+          "Data version mismatch;", data_version, "!=", self.DATA_VERSION
+        )
+      author_id = int.from_bytes(binary[1:9], 'big')
+      targetchannel_id = int.from_bytes(binary[9:17], 'big')
       # from_bytes won't take a single byte, so this hack is needed.
-      self.channeltype_flags = int.from_bytes(bytes((binary[16],)), 'big')
-      reference_id = int.from_bytes(binary[17:25], 'big')
+      self.channeltype_flags = int.from_bytes(bytes((binary[17],)), 'big')
+      reference_id = int.from_bytes(binary[18:26], 'big')
     else:
-      raise CorruptConfessionDataException()
+      raise CorruptConfessionDataException("Data format incorrect;", len(binary), "!=", 26)
     self.author = await self.bot.fetch_user(author_id)
     self.targetchannel = await self.bot.fetch_channel(targetchannel_id)
     self.anonid = self.get_anonid(self.targetchannel.guild.id, self.author.id)
@@ -553,9 +565,10 @@ async def add_image(self, *, attachment:discord.Attachment | None = None, url:st
   def store(self) -> str:
     """ Encrypt data for secure storage """
     # Size limit: ~100 bytes
+    bversion = self.DATA_VERSION.to_bytes(1, 'big')
     bauthor = self.author.id.to_bytes(8, 'big')
     btarget = self.targetchannel.id.to_bytes(8, 'big')
-    bmarket = self.channeltype_flags.to_bytes(1, 'big')
+    bflags = self.channeltype_flags.to_bytes(1, 'big')
     if self.reference:
       breference = self.reference.id.to_bytes(8, 'big')
       # Store in cache so it can be restored
@@ -568,18 +581,21 @@ def store(self) -> str:
     else:
       breference = int(0).to_bytes(8, 'big')
 
-    binary = bauthor + btarget + bmarket + breference
-    return self.parent.crypto.encrypt(binary).decode('ascii')
+    binary = bversion + bauthor + btarget + bflags + breference
+    return b64encode(self.parent.crypto.encrypt(binary)).decode('ascii')
 
   # Data rehydration
 
   def get_anonid(self, guildid:int, userid:int) -> str:
     """ Calculates the current anon-id for a user """
-    offset = int(self.config.get(f"{guildid}_shuffle", fallback=0))
-    encrypted = self.parent.crypto.encrypt(
-      guildid.to_bytes(8, 'big') + userid.to_bytes(8, 'big') + offset.to_bytes(2, 'big')
+    salt = b64decode(self.config.get(f"{guildid}_shuffle", fallback=''))
+    if len(salt) < 16: # If server does not yet have a salt
+      salt = self.parent.crypto.srandom_token()
+      self.config[f"{guildid}_shuffle"] = b64encode(salt).decode('ascii')
+    hashed = self.parent.crypto.hash(
+      guildid.to_bytes(8, 'big') + userid.to_bytes(8, 'big'), salt
     )
-    return hashlib.sha256(encrypted).hexdigest()[-6:]
+    return hashed.hex()[-6:]
 
   def generate_embed(self):
     """ Generate or add anonid to the confession embed """

extensions/confessions_marketplace.py

@@ -5,6 +5,7 @@
 
 from typing import Optional, TYPE_CHECKING
 from enum import IntEnum
+from base64 import b64encode, b64decode
 import discord
 from discord import app_commands
 from discord.ext import commands
@@ -164,9 +165,8 @@ async def on_create_offer(self, inter:discord.Interaction):
       await inter.response.send_message(self.babel(inter, 'error_embed_deleted'), ephemeral=True)
       return
     id_seller = inter.data.get('custom_id')[28:]
-    id_buyer = (
-      self.bot.cogs['Confessions'].crypto.encrypt(inter.user.id.to_bytes(8, 'big')).decode('ascii')
-    )
+    raw_buyer = self.bot.cogs['Confessions'].crypto.encrypt(inter.user.id.to_bytes(8, 'big'))
+    id_buyer = b64encode(raw_buyer).decode('ascii')
     if id_seller == id_buyer:
       await inter.response.send_message(self.babel(inter, 'error_self_offer'), ephemeral=True)
       return
@@ -182,8 +182,10 @@ async def on_accept_offer(self, inter:discord.Interaction):
       return
     encrypted_data = inter.data.get('custom_id')[29:].split('_')
 
-    seller_id = int.from_bytes(self.bot.cogs['Confessions'].crypto.decrypt(encrypted_data[0]), 'big')
-    buyer_id = int.from_bytes(self.bot.cogs['Confessions'].crypto.decrypt(encrypted_data[1]), 'big')
+    raw_seller_id = self.bot.cogs['Confessions'].crypto.decrypt(encrypted_data[0])
+    seller_id = int.from_bytes(b64decode(raw_seller_id), 'big')
+    raw_buyer_id = self.bot.cogs['Confessions'].crypto.decrypt(encrypted_data[1])
+    buyer_id = int.from_bytes(b64decode(raw_buyer_id), 'big')
     if seller_id == inter.user.id:
       seller = inter.user
       buyer = await inter.guild.fetch_member(buyer_id)
@@ -210,7 +212,8 @@ async def on_accept_offer(self, inter:discord.Interaction):
 
   async def on_withdraw(self, inter:discord.Interaction):
     encrypted_data = inter.data.get('custom_id')[31:].split('_')
-    owner_id = int.from_bytes(self.bot.cogs['Confessions'].crypto.decrypt(encrypted_data[-1]), 'big')
+    raw_owner_id = self.bot.cogs['Confessions'].crypto.decrypt(encrypted_data[-1])
+    owner_id = int.from_bytes(b64decode(raw_owner_id), 'big')
     if owner_id != inter.user.id:
       await inter.response.send_message(
         self.babel(inter, 'error_wrong_person', buy=False), ephemeral=True
@@ -290,15 +293,17 @@ async def on_channeltype_send(
   ) -> dict[str] | bool:
     """ Add a view for headed for a marketplace channnel """
     if data.channeltype_flags == MarketplaceFlags.LISTING:
-      id_seller = data.parent.crypto.encrypt(data.author.id.to_bytes(8, 'big')).decode('ascii')
+      raw_seller = data.parent.crypto.encrypt(data.author.id.to_bytes(8, 'big'))
+      id_seller = b64encode(raw_seller).decode('ascii')
       return {
         'use_webhook': False,
         'view': self.ListingView(self, inter, id_seller)
       }
     elif data.channeltype_flags == MarketplaceFlags.OFFER:
       listing = await data.targetchannel.fetch_message(data.reference.id)
       id_seller = listing.components[0].children[0].custom_id[28:]
-      id_buyer = data.parent.crypto.encrypt(data.author.id.to_bytes(8, 'big')).decode('ascii')
+      raw_buyer = data.parent.crypto.encrypt(data.author.id.to_bytes(8, 'big'))
+      id_buyer = b64encode(raw_buyer).decode('ascii')
       return {
         'use_webhook': False,
         'view': self.OfferView(self, inter, id_seller, id_buyer)

extensions/confessions_setup.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import asyncio
+from base64 import b64encode
 from typing import TYPE_CHECKING
 import discord
 from discord import app_commands
@@ -396,8 +397,8 @@ async def shuffle(self, inter:discord.Interaction):
     await inter.response.send_message(self.babel(inter, 'shufflesuccess'))
 
   def perform_shuffle(self, guild_id:int):
-    shuffle = int(self.config.get(str(guild_id) + '_shuffle', fallback=0))
-    self.bot.config.set(self.SCOPE, str(guild_id) + '_shuffle', str(shuffle + 1))
+    salt = self.parent.crypto.srandom_token()
+    self.bot.config.set(self.SCOPE, str(guild_id) + '_shuffle', b64encode(salt).decode('ascii'))
     self.bot.config.save()
 
 

requirements.txt

@@ -1,2 +1,2 @@
 discord.py==2.4.0
-cryptography
+pycryptodome