encrypt.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170
  1. #!/usr/bin/env python
  2. # -*- coding: iso-8859-15 -*-
  3. """
  4. $Id$
  5. This file is part of the anontwi project, http://anontwi.03c8.net
  6. Copyright (c) 2012/2013/2014/2015 by psy <epsylon@riseup.net>
  7. anontwi is free software; you can redistribute it and/or modify it under
  8. the terms of the GNU General Public License as published by the Free
  9. Software Foundation version 3 of the License.
  10. anontwi is distributed in the hope that it will be useful, but WITHOUT ANY
  11. WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  12. FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
  13. details.
  14. You should have received a copy of the GNU General Public License along
  15. with anontwi; if not, write to the Free Software Foundation, Inc., 51
  16. Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
  17. """
  18. ###################################################################
  19. # See https://en.wikipedia.org/wiki/HMAC#Implementation
  20. # Example written by: michael@briarproject.org
  21. ###################################################################
  22. # Constants for AES256 and HMAC-SHA1
  23. KEY_SIZE = 32
  24. BLOCK_SIZE = 16
  25. MAC_SIZE = 20
  26. from os import urandom
  27. from hashlib import sha1, sha256
  28. from Crypto.Cipher import AES
  29. from base64 import b64encode, b64decode
  30. trans_5C = "".join([chr (x ^ 0x5c) for x in xrange(256)])
  31. trans_36 = "".join([chr (x ^ 0x36) for x in xrange(256)])
  32. def hmac_sha1(key, msg):
  33. if len(key) > 20:
  34. key = sha1(key).digest()
  35. key += chr(0) * (20 - len(key))
  36. o_key_pad = key.translate(trans_5C)
  37. i_key_pad = key.translate(trans_36)
  38. return sha1(o_key_pad + sha1(i_key_pad + msg).digest()).digest()
  39. def derive_keys(key):
  40. h = sha256()
  41. h.update(key)
  42. h.update('cipher')
  43. cipher_key = h.digest()
  44. h = sha256()
  45. h.update(key)
  46. h.update('mac')
  47. mac_key = h.digest()
  48. return (cipher_key, mac_key)
  49. def generate_key():
  50. return b64encode(urandom(KEY_SIZE))
  51. class Cipher(object):
  52. """
  53. Cipher class
  54. """
  55. def __init__(self, key="", text=""):
  56. """
  57. Init
  58. """
  59. self.block_size = 16
  60. self.mac_size = 20
  61. self.key = self.set_key(key)
  62. self.text = self.set_text(text)
  63. self.mode = AES.MODE_CFB
  64. def set_key(self, key):
  65. """
  66. Set key
  67. """
  68. # Base64 decode the key
  69. try:
  70. key = b64decode(key)
  71. except TypeError:
  72. raise ValueError
  73. # The key must be the expected length
  74. if len(key) != KEY_SIZE:
  75. raise ValueError
  76. self.key = key
  77. return self.key
  78. def set_text(self, text):
  79. """
  80. Set text
  81. """
  82. self.text = text
  83. return self.text
  84. def encrypt(self):
  85. """
  86. Encrypt text
  87. """
  88. # The IV, ciphertext and MAC can't be more than 105 bytes
  89. if BLOCK_SIZE + len(self.text) + MAC_SIZE > 105:
  90. self.text = self.text[:105 - BLOCK_SIZE - MAC_SIZE]
  91. # Derive the cipher and MAC keys
  92. (cipher_key, mac_key) = derive_keys(self.key)
  93. # Generate a random IV
  94. iv = urandom(BLOCK_SIZE)
  95. # Encrypt the plaintext
  96. aes = AES.new(cipher_key, self.mode, iv)
  97. ciphertext = aes.encrypt(self.text)
  98. # Calculate the MAC over the IV and the ciphertext
  99. mac = hmac_sha1(mac_key, iv + ciphertext)
  100. # Base64 encode the IV, ciphertext and MAC
  101. return b64encode(iv + ciphertext + mac)
  102. def decrypt(self):
  103. """
  104. Decrypt text
  105. """
  106. # Base64 decode
  107. try:
  108. iv_ciphertext_mac = b64decode(self.text)
  109. except TypeError:
  110. return None
  111. # Separate the IV, ciphertext and MAC
  112. iv = iv_ciphertext_mac[:BLOCK_SIZE]
  113. ciphertext = iv_ciphertext_mac[BLOCK_SIZE:-MAC_SIZE]
  114. mac = iv_ciphertext_mac[-MAC_SIZE:]
  115. # Derive the cipher and MAC keys
  116. (cipher_key, mac_key) = derive_keys(self.key)
  117. # Calculate the expected MAC
  118. expected_mac = hmac_sha1(mac_key, iv + ciphertext)
  119. # Check the MAC
  120. if mac != expected_mac:
  121. return None
  122. # Decrypt the ciphertext
  123. aes = AES.new(cipher_key, self.mode, iv)
  124. return aes.decrypt(ciphertext)
  125. if __name__ == "__main__":
  126. key = generate_key()
  127. print 'Key:', key
  128. # Encrypt and decrypt a short message
  129. text = 'Hello world!'
  130. c = Cipher(key, text)
  131. msg = c.encrypt()
  132. c.set_text(msg)
  133. print '\nCiphertext:', msg
  134. print 'Length:', len(msg)
  135. print 'Plaintext:', c.decrypt()
  136. # Encrypt and decrypt a long message
  137. text = 'Gosh this is a long message, far too long to fit in a tweet I dare say, especially when you consider the encryption overhead'
  138. c = Cipher(key, text)
  139. msg = c.encrypt()
  140. c.set_text(msg)
  141. print '\nCiphertext:', msg
  142. print 'Length:', len(msg)
  143. print 'Plaintext:', c.decrypt()
  144. # Check that modifying the message invalidates the MAC
  145. text = 'Hello world!'
  146. c = Cipher(key, text)
  147. msg = c.encrypt()
  148. msg = msg[:16] + msg[17] + msg[16] + msg[18:]
  149. c.set_text(msg)
  150. print '\nCiphertext:', msg
  151. print 'Length:', len(msg)
  152. print 'Plaintext:', c.decrypt()