encodehelpers.py 8.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154
  1. from __future__ import print_function
  2. import base64
  3. import codecs
  4. import sys
  5. from .compat import isnumber
  6. STRING_ESCAPE_RAW = 'raw'
  7. STRING_ESCAPE_PRINT = 'print'
  8. STRING_ESCAPE_UTF8 = 'utf8'
  9. STRING_ESCAPE_BASE64 = 'base64'
  10. ESCAPE_CHOICES = [STRING_ESCAPE_RAW, STRING_ESCAPE_PRINT, STRING_ESCAPE_UTF8, STRING_ESCAPE_BASE64]
  11. if sys.version_info < (3,):
  12. bval = ord
  13. def num2unistr(i): return codecs.decode(str(i), 'ascii')
  14. num2bytes = str
  15. else:
  16. def bval(x): return x
  17. num2unistr = str
  18. def num2bytes(i): return codecs.encode(str(i), 'ascii')
  19. ASCII_ESCAPE_LOOKUP = [u'\\x00', u'\\x01', u'\\x02', u'\\x03', u'\\x04', u'\\x05', u'\\x06', u'\\x07', u'\\x08',
  20. u'\\x09', u'\\x0A', u'\\x0B', u'\\x0C', u'\\x0D', u'\\x0E', u'\\x0F', u'\\x10', u'\\x11',
  21. u'\\x12', u'\\x13', u'\\x14', u'\\x15', u'\\x16', u'\\x17', u'\\x18', u'\\x19', u'\\x1A',
  22. u'\\x1B', u'\\x1C', u'\\x1D', u'\\x1E', u'\\x1F', u' ', u'!', u'"', u'#', u'$', u'%', u'&', u"'",
  23. u'(', u')', u'*', u'+', u',', u'-', u'.', u'/', u'0', u'1', u'2', u'3', u'4', u'5', u'6', u'7',
  24. u'8', u'9', u':', u';', u'<', u'=', u'>', u'?', u'@', u'A', u'B', u'C', u'D', u'E', u'F', u'G',
  25. u'H', u'I', u'J', u'K', u'L', u'M', u'N', u'O', u'P', u'Q', u'R', u'S', u'T', u'U', u'V', u'W',
  26. u'X', u'Y', u'Z', u'[', u'\\', u']', u'^', u'_', u'`', u'a', u'b', u'c', u'd', u'e', u'f', u'g',
  27. u'h', u'i', u'j', u'k', u'l', u'm', u'n', u'o', u'p', u'q', u'r', u's', u't', u'u', u'v', u'w',
  28. u'x', u'y', u'z', u'{', u'|', u'}', u'~', u'\\x7F', u'\\x80', u'\\x81', u'\\x82', u'\\x83',
  29. u'\\x84', u'\\x85', u'\\x86', u'\\x87', u'\\x88', u'\\x89', u'\\x8A', u'\\x8B', u'\\x8C',
  30. u'\\x8D', u'\\x8E', u'\\x8F', u'\\x90', u'\\x91', u'\\x92', u'\\x93', u'\\x94', u'\\x95',
  31. u'\\x96', u'\\x97', u'\\x98', u'\\x99', u'\\x9A', u'\\x9B', u'\\x9C', u'\\x9D', u'\\x9E',
  32. u'\\x9F', u'\\xA0', u'\\xA1', u'\\xA2', u'\\xA3', u'\\xA4', u'\\xA5', u'\\xA6', u'\\xA7',
  33. u'\\xA8', u'\\xA9', u'\\xAA', u'\\xAB', u'\\xAC', u'\\xAD', u'\\xAE', u'\\xAF', u'\\xB0',
  34. u'\\xB1', u'\\xB2', u'\\xB3', u'\\xB4', u'\\xB5', u'\\xB6', u'\\xB7', u'\\xB8', u'\\xB9',
  35. u'\\xBA', u'\\xBB', u'\\xBC', u'\\xBD', u'\\xBE', u'\\xBF', u'\\xC0', u'\\xC1', u'\\xC2',
  36. u'\\xC3', u'\\xC4', u'\\xC5', u'\\xC6', u'\\xC7', u'\\xC8', u'\\xC9', u'\\xCA', u'\\xCB',
  37. u'\\xCC', u'\\xCD', u'\\xCE', u'\\xCF', u'\\xD0', u'\\xD1', u'\\xD2', u'\\xD3', u'\\xD4',
  38. u'\\xD5', u'\\xD6', u'\\xD7', u'\\xD8', u'\\xD9', u'\\xDA', u'\\xDB', u'\\xDC', u'\\xDD',
  39. u'\\xDE', u'\\xDF', u'\\xE0', u'\\xE1', u'\\xE2', u'\\xE3', u'\\xE4', u'\\xE5', u'\\xE6',
  40. u'\\xE7', u'\\xE8', u'\\xE9', u'\\xEA', u'\\xEB', u'\\xEC', u'\\xED', u'\\xEE', u'\\xEF',
  41. u'\\xF0', u'\\xF1', u'\\xF2', u'\\xF3', u'\\xF4', u'\\xF5', u'\\xF6', u'\\xF7', u'\\xF8',
  42. u'\\xF9', u'\\xFA', u'\\xFB', u'\\xFC', u'\\xFD', u'\\xFE', u'\\xFF']
  43. ASCII_ESCAPE_LOOKUP_BYTES = [b'\\x00', b'\\x01', b'\\x02', b'\\x03', b'\\x04', b'\\x05', b'\\x06', b'\\x07', b'\\x08',
  44. b'\\x09', b'\\x0A', b'\\x0B', b'\\x0C', b'\\x0D', b'\\x0E', b'\\x0F', b'\\x10', b'\\x11',
  45. b'\\x12', b'\\x13', b'\\x14', b'\\x15', b'\\x16', b'\\x17', b'\\x18', b'\\x19', b'\\x1A',
  46. b'\\x1B', b'\\x1C', b'\\x1D', b'\\x1E', b'\\x1F', b' ', b'!', b'"', b'#', b'$', b'%', b'&',
  47. b"'", b'(', b')', b'*', b'+', b',', b'-', b'.', b'/', b'0', b'1', b'2', b'3', b'4', b'5',
  48. b'6', b'7', b'8', b'9', b':', b';', b'<', b'=', b'>', b'?', b'@', b'A', b'B', b'C', b'D',
  49. b'E', b'F', b'G', b'H', b'I', b'J', b'K', b'L', b'M', b'N', b'O', b'P', b'Q', b'R', b'S',
  50. b'T', b'U', b'V', b'W', b'X', b'Y', b'Z', b'[', b'\\', b']', b'^', b'_', b'`', b'a', b'b',
  51. b'c', b'd', b'e', b'f', b'g', b'h', b'i', b'j', b'k', b'l', b'm', b'n', b'o', b'p', b'q',
  52. b'r', b's', b't', b'u', b'v', b'w', b'x', b'y', b'z', b'{', b'|', b'}', b'~', b'\\x7F',
  53. b'\\x80', b'\\x81', b'\\x82', b'\\x83', b'\\x84', b'\\x85', b'\\x86', b'\\x87', b'\\x88',
  54. b'\\x89', b'\\x8A', b'\\x8B', b'\\x8C', b'\\x8D', b'\\x8E', b'\\x8F', b'\\x90', b'\\x91',
  55. b'\\x92', b'\\x93', b'\\x94', b'\\x95', b'\\x96', b'\\x97', b'\\x98', b'\\x99', b'\\x9A',
  56. b'\\x9B', b'\\x9C', b'\\x9D', b'\\x9E', b'\\x9F', b'\\xA0', b'\\xA1', b'\\xA2', b'\\xA3',
  57. b'\\xA4', b'\\xA5', b'\\xA6', b'\\xA7', b'\\xA8', b'\\xA9', b'\\xAA', b'\\xAB', b'\\xAC',
  58. b'\\xAD', b'\\xAE', b'\\xAF', b'\\xB0', b'\\xB1', b'\\xB2', b'\\xB3', b'\\xB4', b'\\xB5',
  59. b'\\xB6', b'\\xB7', b'\\xB8', b'\\xB9', b'\\xBA', b'\\xBB', b'\\xBC', b'\\xBD', b'\\xBE',
  60. b'\\xBF', b'\\xC0', b'\\xC1', b'\\xC2', b'\\xC3', b'\\xC4', b'\\xC5', b'\\xC6', b'\\xC7',
  61. b'\\xC8', b'\\xC9', b'\\xCA', b'\\xCB', b'\\xCC', b'\\xCD', b'\\xCE', b'\\xCF', b'\\xD0',
  62. b'\\xD1', b'\\xD2', b'\\xD3', b'\\xD4', b'\\xD5', b'\\xD6', b'\\xD7', b'\\xD8', b'\\xD9',
  63. b'\\xDA', b'\\xDB', b'\\xDC', b'\\xDD', b'\\xDE', b'\\xDF', b'\\xE0', b'\\xE1', b'\\xE2',
  64. b'\\xE3', b'\\xE4', b'\\xE5', b'\\xE6', b'\\xE7', b'\\xE8', b'\\xE9', b'\\xEA', b'\\xEB',
  65. b'\\xEC', b'\\xED', b'\\xEE', b'\\xEF', b'\\xF0', b'\\xF1', b'\\xF2', b'\\xF3', b'\\xF4',
  66. b'\\xF5', b'\\xF6', b'\\xF7', b'\\xF8', b'\\xF9', b'\\xFA', b'\\xFB', b'\\xFC', b'\\xFD',
  67. b'\\xFE', b'\\xFF']
  68. def escape_ascii(bytes_data):
  69. return u''.join(ASCII_ESCAPE_LOOKUP[bval(ch)] for ch in bytes_data)
  70. def escape_ascii_bytes(bytes_data):
  71. return b''.join(ASCII_ESCAPE_LOOKUP_BYTES[bval(ch)] for ch in bytes_data)
  72. def escape_utf8_error(err):
  73. return escape_ascii(err.object[err.start:err.end]), err.end
  74. codecs.register_error('rdbslashescape', escape_utf8_error)
  75. def escape_utf8(byte_data):
  76. return byte_data.decode('utf-8', 'rdbslashescape')
  77. def bytes_to_unicode(byte_data, escape, skip_printable=False):
  78. """
  79. Decode given bytes using specified escaping method.
  80. :param byte_data: The byte-like object with bytes to decode.
  81. :param escape: The escape method to use.
  82. :param skip_printable: If True, don't escape byte_data with all 'printable ASCII' bytes. Defaults to False.
  83. :return: New unicode string, escaped with the specified method if needed.
  84. """
  85. if isnumber(byte_data):
  86. if skip_printable:
  87. return num2unistr(byte_data)
  88. else:
  89. byte_data = num2bytes(byte_data)
  90. else:
  91. assert (isinstance(byte_data, type(b'')))
  92. if skip_printable and all(0x20 <= bval(ch) <= 0x7E for ch in byte_data):
  93. escape = STRING_ESCAPE_RAW
  94. if escape == STRING_ESCAPE_RAW:
  95. return byte_data.decode('latin-1')
  96. elif escape == STRING_ESCAPE_PRINT:
  97. return escape_ascii(byte_data)
  98. elif escape == STRING_ESCAPE_UTF8:
  99. return escape_utf8(byte_data)
  100. elif escape == STRING_ESCAPE_BASE64:
  101. return codecs.decode(base64.b64encode(byte_data), 'latin-1')
  102. else:
  103. raise UnicodeEncodeError("Unknown escape option")
  104. def apply_escape_bytes(byte_data, escape, skip_printable=False):
  105. """
  106. Apply the specified escape method on the given bytes.
  107. :param byte_data: The byte-like object with bytes to escape.
  108. :param escape: The escape method to use.
  109. :param skip_printable: If True, don't escape byte_data with all 'printable ASCII' bytes. Defaults to False.
  110. :return: new bytes object with the escaped bytes or byte_data itself on some no-op cases.
  111. """
  112. if isnumber(byte_data):
  113. if skip_printable:
  114. return num2bytes(byte_data)
  115. else:
  116. byte_data = num2bytes(byte_data)
  117. else:
  118. assert (isinstance(byte_data, type(b'')))
  119. if skip_printable and all(0x20 <= bval(ch) <= 0x7E for ch in byte_data):
  120. escape = STRING_ESCAPE_RAW
  121. if escape == STRING_ESCAPE_RAW:
  122. return byte_data
  123. elif escape == STRING_ESCAPE_PRINT:
  124. return escape_ascii_bytes(byte_data)
  125. elif escape == STRING_ESCAPE_UTF8:
  126. return codecs.encode(escape_utf8(byte_data), 'utf-8')
  127. elif escape == STRING_ESCAPE_BASE64:
  128. return base64.b64encode(byte_data)
  129. else:
  130. raise UnicodeEncodeError("Unknown escape option")