crypto_message.py 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157
  1. # The contents of this file are subject to the BitTorrent Open Source License
  2. # Version 1.1 (the License). You may not copy or use this file, in either
  3. # source code or executable form, except in compliance with the License. You
  4. # may obtain a copy of the License at http://www.bittorrent.com/license/.
  5. #
  6. # Software distributed under the License is distributed on an AS IS basis,
  7. # WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
  8. # for the specific language governing rights and limitations under the
  9. # License.
  10. # by Benjamin C. Wiley Sittler
  11. import Crypto.Cipher.AES as _AES
  12. import sha as _sha
  13. import os as _os
  14. import hmac as _hmac
  15. import string as _string
  16. _urlbase64 = _string.maketrans('+/-_', '-_+/')
  17. def pad(data, length):
  18. '''
  19. PKCS #7-style padding with the given block length
  20. '''
  21. assert length < 256
  22. assert length > 0
  23. padlen = length - len(data) % length
  24. assert padlen <= length
  25. assert padlen > 0
  26. return data + padlen * chr(padlen)
  27. def unpad(data, length):
  28. '''
  29. PKCS #7-style unpadding with the given block length
  30. '''
  31. assert length < 256
  32. assert length > 0
  33. padlen = ord(data[-1])
  34. assert padlen <= length
  35. assert padlen > 0
  36. assert data[-padlen:] == padlen * chr(padlen)
  37. return data[:-padlen]
  38. def ascii(data):
  39. '''
  40. Encode data as URL-safe variant of Base64.
  41. '''
  42. return data.encode('base64').translate(_urlbase64, '\r\n')
  43. def unascii(data):
  44. '''
  45. Decode data from URL-safe variant of Base64.
  46. '''
  47. decoded = data.translate(_urlbase64).decode('base64')
  48. assert ascii(decoded) == data
  49. return decoded
  50. def encode(data, secret, salt = None):
  51. '''
  52. Encode and return the data as a random-IV-prefixed AES-encrypted
  53. HMAC-SHA1-authenticated padded message corresponding to the given
  54. data string and secret, which should be at least 36 randomly
  55. chosen bytes agreed upon by the encoding and decoding parties.
  56. '''
  57. assert len(secret) >= 36
  58. if salt is None:
  59. salt = _os.urandom(16)
  60. aes = _AES.new(secret[:16], _AES.MODE_CBC, salt)
  61. padded_data = pad(20 * '\0' + data, 16)[20:]
  62. mac = _hmac.HMAC(key = secret[16:], msg = padded_data, digestmod = _sha).digest()
  63. encrypted = aes.encrypt(mac + padded_data)
  64. return salt + encrypted
  65. def decode(data, secret):
  66. '''
  67. Decode and return the data from random-IV-prefixed AES-encrypted
  68. HMAC-SHA1-authenticated padded message corresponding to the given
  69. data string and secret, which should be at least 36 randomly
  70. chosen bytes agreed upon by the encoding and decoding parties.
  71. '''
  72. assert len(secret) >= 36
  73. salt = data[:16]
  74. encrypted = data[16:]
  75. aes = _AES.new(secret[:16], _AES.MODE_CBC, salt)
  76. decrypted = aes.decrypt(encrypted)
  77. mac = decrypted[:20]
  78. padded_data = decrypted[20:]
  79. mac2 = _hmac.HMAC(key = secret[16:], msg = padded_data, digestmod = _sha).digest()
  80. assert mac == mac2
  81. return unpad(20 * '\0' + padded_data, 16)[20:]
  82. def test():
  83. '''
  84. Trivial smoke test to make sure this module works.
  85. '''
  86. secret = unascii('D_4j_P5Fh-UWUuH2U3IYw2erxRab5QX0zOR7eYlucT0GfuuwxgoGcfKI_rnyStbllZTPBbCESbKv0kMsUB9tOnLvAU2k7bCcMy7ylUqFwgc=')
  87. secret2 = unascii('e3YUIIA3APP66cMJrKNRAHVm0nd7BRAxZqyiYadTML78v2yS')
  88. salt = unascii('yRja3Cj5qc2xhYoSJtCBSw==')
  89. for data, message, message2 in (
  90. ('Hello, world!',
  91. 'yRja3Cj5qc2xhYoSJtCBSxqHihP8mZ8TNuiLv_i41uaHM8jUu4N2cpU_XmlH0raoq-6FLOHE3ScV9aPnQ9Ulsg==',
  92. 'yRja3Cj5qc2xhYoSJtCBS8MPPvak9ZDXydyMlACoQ7WSlM7X4PunKhJa775itirxJPD1eFgSnWHjAjmZn_8bvg==',
  93. ),
  94. ('',
  95. 'yRja3Cj5qc2xhYoSJtCBS6vWZ3nvvsp3gM2-G-co6fVCvkLv6pRrfLQg2vm1yNzr',
  96. 'yRja3Cj5qc2xhYoSJtCBSy9XX0E8Re0XumS1wMMEJFwSkTIQBGqbWGH4_GPMwdrR',
  97. ),
  98. ('\0',
  99. 'yRja3Cj5qc2xhYoSJtCBSyEz2FFkaC3bRhMV03csag5MMIrVaWeWK2J1IXIaK_UQ',
  100. 'yRja3Cj5qc2xhYoSJtCBS-05SxrZqgT9XhcEWp0eTLCrdQpnzBGKLL8qvIsc6nx6',
  101. ),
  102. ('Hi there!',
  103. 'yRja3Cj5qc2xhYoSJtCBS8oy34UlBkk3v__LUHTa557U04HT_-M80DunhcKbFh-q',
  104. 'yRja3Cj5qc2xhYoSJtCBS-6M4ylGA0jmaPjWRiEoBy3j1R1o17_KbsAH_0CiZRhx',
  105. ),
  106. ):
  107. assert unascii(ascii(data)) == data
  108. assert ascii(unascii(message)) == message
  109. assert len(pad(data, 16)) % 16 == 0
  110. assert unpad(pad(data, 16), 16) == data
  111. assert message == ascii(encode(data, secret, salt))
  112. assert decode(unascii(message), secret) == data
  113. assert decode(encode(data, secret), secret) == data
  114. assert message2 == ascii(encode(data, secret2, salt))
  115. assert decode(unascii(message2), secret2) == data
  116. assert decode(encode(data, secret2), secret2) == data
  117. test()
  118. def main(sys):
  119. progname = sys.argv[0]
  120. secret = _os.urandom(36)
  121. salt = None
  122. if len(sys.argv) < 2:
  123. sys.stderr.write('%s: secret is %s\n' % (progname, ascii(secret)))
  124. sys.stderr.flush()
  125. elif len(sys.argv) < 3:
  126. progname, secret = sys.argv
  127. secret = unascii(secret)
  128. else:
  129. progname, secret, salt = sys.argv
  130. secret = unascii(secret)
  131. salt = unascii(salt)
  132. while True:
  133. line = sys.stdin.readline()
  134. if not line:
  135. break
  136. try:
  137. sys.stdout.write('%s' % decode(unascii(line.rstrip('\r\n')), secret))
  138. except:
  139. sys.stdout.write('%s\n' % ascii(encode(line, secret, salt)))
  140. sys.stdout.flush()
  141. if __name__ == '__main__':
  142. import sys
  143. main(sys)