ctypescrypto/cipher.py

   1 from ctypes import *\r
   2 \r
   3 CIPHER_ALGORITHMS = ("DES", "DES-EDE3", "BF", "AES-128", "AES-192", "AES-256")\r
   4 CIPHER_MODES = ("CBC", "CFB", "OFB", "ECB")\r
   5 \r
   6 class CipherError(Exception):\r
   7     pass\r
   8 \r
   9 class CipherType:\r
  10 \r
  11     def __init__(self, libcrypto, cipher_algo, cipher_mode):\r
  12         self.libcrypto = libcrypto\r
  13         self.cipher_algo = cipher_algo\r
  14         self.cipher_mode = cipher_mode\r
  15         cipher_name = "-".join([self.cipher_algo, self.cipher_mode])\r
  16         self.cipher = self.libcrypto.EVP_get_cipherbyname(cipher_name)\r
  17         if self.cipher == 0:\r
  18             raise CipherError, "Unknown cipher: %s" % cipher_name\r
  19 \r
  20     def __del__(self):\r
  21         pass\r
  22 \r
  23     def algo(self):\r
  24         return self.cipher_algo\r
  25 \r
  26     def mode(self):\r
  27         return self.cipher_mode\r
  28 \r
  29 class Cipher:\r
  30 \r
  31     def __init__(self, libcrypto, cipher_type, key, iv, encrypt=True):\r
  32         self.libcrypto = libcrypto\r
  33         self._clean_ctx()\r
  34         key_ptr = c_char_p(key)\r
  35         iv_ptr = c_char_p(iv)\r
  36         self.ctx = self.libcrypto.EVP_CIPHER_CTX_new(cipher_type.cipher, None, key_ptr, iv_ptr)\r
  37         if self.ctx == 0:\r
  38             raise CipherError, "Unable to create cipher context"\r
  39         self.encrypt = encrypt\r
  40         if encrypt: \r
  41             enc = 1\r
  42         else: \r
  43             enc = 0\r
  44         result = self.libcrypto.EVP_CipherInit_ex(self.ctx, cipher_type.cipher, None, key_ptr, iv_ptr, c_int(enc))\r
  45         self.cipher_type = cipher_type\r
  46         if result == 0:\r
  47             self._clean_ctx()\r
  48             raise CipherError, "Unable to initialize cipher"\r
  49 \r
  50     def __del__(self):\r
  51         self._clean_ctx()\r
  52 \r
  53     def enable_padding(self, padding=True):\r
  54         if padding:\r
  55             padding_flag = 1\r
  56         else:\r
  57             padding_flag = 0\r
  58         self.libcrypto.EVP_CIPHER_CTX_set_padding(self.ctx, padding_flag)\r
  59 \r
  60     def update(self, data):\r
  61         if self.cipher_finalized :\r
  62             raise CipherError, "No updates allowed"\r
  63         if type(data) != type(""):\r
  64             raise TypeError, "A string is expected"\r
  65         if len(data) <= 0:\r
  66             return ""\r
  67         self.data = self.data + data\r
  68     \r
  69     def finish(self, data=None):\r
  70         if data is not None:\r
  71             self.update(data)\r
  72         return self._finish()\r
  73         \r
  74     def _finish(self):\r
  75         if self.cipher_finalized :\r
  76             raise CipherError, "Cipher operation is already completed"\r
  77         self.cipher_out = create_string_buffer(len(self.data) + 32)\r
  78         result = self.libcrypto.EVP_CipherUpdate(self.ctx, byref(self.cipher_out), byref(self.cipher_out_len), c_char_p(self.data), len(self.data))\r
  79         if result == 0:\r
  80             self._clean_ctx()\r
  81             raise CipherError, "Unable to update cipher"\r
  82         self.cipher_finalized = True\r
  83         update_data = self.cipher_out.raw[:self.cipher_out_len.value]\r
  84         result = self.libcrypto.EVP_CipherFinal_ex(self.ctx, byref(self.cipher_out), byref(self.cipher_out_len))\r
  85         if result == 0:\r
  86             self._clean_ctx()\r
  87             raise CipherError, "Unable to finalize cipher"\r
  88         final_data = self.cipher_out.raw[:self.cipher_out_len.value]\r
  89         return update_data + final_data\r
  90         \r
  91     def _clean_ctx(self):\r
  92         try:\r
  93             if self.ctx is not None:\r
  94                 self.libcrypto.EVP_CIPHER_CTX_cleanup(self.ctx)\r
  95                 self.libcrypto.EVP_CIPHER_CTX_free(self.ctx)\r
  96                 del(self.ctx)\r
  97         except AttributeError:\r
  98             pass\r
  99         self.cipher_out = None\r
 100         self.cipher_out_len = c_long(0)\r
 101         self.data = ""\r
 102         self.cipher_finalized = False