crypto.py 2.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667
  1. import os
  2. import zlib
  3. from Crypto.Cipher import AES
  4. from Crypto.Hash import SHA256, HMAC
  5. from Crypto.Util.number import bytes_to_long, long_to_bytes
  6. from .helpers import IntegrityError
  7. from .oaep import OAEP
  8. class CryptoManager(object):
  9. CREATE = '\1'
  10. READ = '\2'
  11. def __init__(self, store):
  12. self.key_cache = {}
  13. self.store = store
  14. self.tid = store.tid
  15. self.id_key = '0' * 32
  16. self.read_key = os.urandom(32)
  17. self.create_key = os.urandom(32)
  18. self.read_encrypted = OAEP(256, hash=SHA256).encode(self.read_key, os.urandom(32))
  19. self.create_encrypted = OAEP(256, hash=SHA256).encode(self.create_key, os.urandom(32))
  20. def id_hash(self, data):
  21. return HMAC.new(self.id_key, data, SHA256).digest()
  22. def encrypt_read(self, data):
  23. key_data = OAEP(256, hash=SHA256).encode(self.read_key, os.urandom(32))
  24. #key_data = self.rsa_create.encrypt(key_data)
  25. data = zlib.compress(data)
  26. hash = SHA256.new(data).digest()
  27. data = AES.new(self.read_key, AES.MODE_CFB, hash[:16]).encrypt(data)
  28. return ''.join((self.READ, self.read_encrypted, hash, data))
  29. def encrypt_create(self, data):
  30. key_data = OAEP(256, hash=SHA256).encode(self.create_key, os.urandom(32))
  31. #key_data = self.rsa_create.encrypt(key_data)
  32. data = zlib.compress(data)
  33. hash = SHA256.new(data).digest()
  34. data = AES.new(self.create_key, AES.MODE_CFB, hash[:16]).encrypt(data)
  35. return ''.join((self.CREATE, self.create_encrypted, hash, data))
  36. def decrypt(self, data):
  37. type = data[0]
  38. if type == self.READ:
  39. key_data = data[1:257]
  40. hash = data[257:289]
  41. #key_data = self.rsa_create.decrypt(key_data)
  42. key = OAEP(256, hash=SHA256).decode(key_data)
  43. data = AES.new(key, AES.MODE_CFB, hash[:16]).decrypt(data[289:])
  44. if SHA256.new(data).digest() != hash:
  45. raise IntegrityError('decryption failed')
  46. return zlib.decompress(data)
  47. elif type == self.CREATE:
  48. key_data = data[1:257]
  49. hash = data[257:289]
  50. #key_data = self.rsa_create.decrypt(key_data)
  51. key = OAEP(256, hash=SHA256).decode(key_data)
  52. data = AES.new(key, AES.MODE_CFB, hash[:16]).decrypt(data[289:])
  53. if SHA256.new(data).digest() != hash:
  54. raise IntegrityError('decryption failed')
  55. return zlib.decompress(data)
  56. else:
  57. raise Exception('Unknown pack type %d found' % ord(type))