misc_ops.py 3.03 KB
Newer Older
1
import builtins
2 3 4
import sys
import hashlib

mathieui's avatar
mathieui committed
5
from typing import Optional, Union, Callable, List
6

mathieui's avatar
mathieui committed
7 8 9 10
bytes_ = builtins.bytes  # alias the stdlib type but ew


def unicode(text: Union[bytes_, str]) -> str:
11
    if not isinstance(text, str):
12 13 14
        return text.decode('utf-8')
    else:
        return text
15 16


mathieui's avatar
mathieui committed
17
def bytes(text: Optional[Union[str, bytes_]]) -> bytes_:
18 19 20 21 22 23 24 25 26
    """
    Convert Unicode text to UTF-8 encoded bytes.

    Since Python 2.6+ and Python 3+ have similar but incompatible
    signatures, this function unifies the two to keep code sane.

    :param text: Unicode text to convert to bytes
    :rtype: bytes (Python3), str (Python2.6+)
    """
Lance Stout's avatar
Lance Stout committed
27 28 29
    if text is None:
        return b''

30 31 32 33 34 35
    if isinstance(text, builtins.bytes):
        # We already have bytes, so do nothing
        return text
    if isinstance(text, list):
        # Convert a list of integers to bytes
        return builtins.bytes(text)
36
    else:
37 38
        # Convert UTF-8 text to bytes
        return builtins.bytes(text, encoding='utf-8')
39 40


mathieui's avatar
mathieui committed
41
def quote(text: Union[str, bytes_]) -> bytes_:
42 43 44 45 46 47 48 49 50
    """
    Enclose in quotes and escape internal slashes and double quotes.

    :param text: A Unicode or byte string.
    """
    text = bytes(text)
    return b'"' + text.replace(b'\\', b'\\\\').replace(b'"', b'\\"') + b'"'


mathieui's avatar
mathieui committed
51
def num_to_bytes(num: int) -> bytes_:
52 53 54 55 56 57 58 59 60 61 62 63 64
    """
    Convert an integer into a four byte sequence.

    :param integer num: An integer to convert to its byte representation.
    """
    bval = b''
    bval += bytes(chr(0xFF & (num >> 24)))
    bval += bytes(chr(0xFF & (num >> 16)))
    bval += bytes(chr(0xFF & (num >> 8)))
    bval += bytes(chr(0xFF & (num >> 0)))
    return bval


mathieui's avatar
mathieui committed
65
def bytes_to_num(bval: bytes_) -> int:
66 67 68 69 70 71
    """
    Convert a four byte sequence to an integer.

    :param bytes bval: A four byte sequence to turn into an integer.
    """
    num = 0
mathieui's avatar
mathieui committed
72 73 74 75
    num += (bval[0] << 24)
    num += (bval[1] << 16)
    num += (bval[2] << 8)
    num += (bval[3])
76 77 78
    return num


mathieui's avatar
mathieui committed
79
def XOR(x: bytes_, y: bytes_) -> bytes_:
80 81 82 83 84 85 86
    """
    Return the results of an XOR operation on two equal length byte strings.

    :param bytes x: A byte string
    :param bytes y: A byte string
    :rtype: bytes
    """
87 88 89
    # This operation is faster with a list comprehension than with a
    # generator, as of 2016 on python 3.5.
    return builtins.bytes([a ^ b for a, b in zip(x, y)])
90 91


mathieui's avatar
mathieui committed
92
def hash(name: str) -> Optional[Callable]:
93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108
    """
    Return a hash function implementing the given algorithm.

    :param name: The name of the hashing algorithm to use.
    :type name: string

    :rtype: function
    """
    name = name.lower()
    if name.startswith('sha-'):
        name = 'sha' + name[4:]
    if name in dir(hashlib):
        return getattr(hashlib, name)
    return None


mathieui's avatar
mathieui committed
109
def hashes() -> List[str]:
110 111 112 113 114 115 116 117 118 119 120 121
    """
    Return a list of available hashing algorithms.

    :rtype: list of strings
    """
    t = []
    if 'md5' in dir(hashlib):
        t = ['MD5']
    if 'md2' in dir(hashlib):
        t += ['MD2']
    hashes = ['SHA-' + h[3:] for h in dir(hashlib) if h.startswith('sha')]
    return t + hashes