import binascii, ctypes, math, struct

class sha2_ctx(ctypes.Structure):
    _fields_ = [("state", ctypes.c_uint32 * 8),
                ("count", ctypes.c_uint64),
                ("buf", ctypes.c_byte * 64)]

sodium = ctypes.CDLL("libsodium.so")

sodium.crypto_hash_sha256_init.argtypes = [ctypes.POINTER(sha2_ctx)]
sodium.crypto_hash_sha256_init.restype = ctypes.c_int

sodium.crypto_hash_sha256_update.argtypes = [ctypes.POINTER(sha2_ctx), ctypes.c_char_p,
                                             ctypes.c_size_t]
sodium.crypto_hash_sha256_update.restype = ctypes.c_int

sodium.crypto_hash_sha256_final.argtypes = [ctypes.POINTER(sha2_ctx), ctypes.c_char_p]
sodium.crypto_hash_sha256_final.restype = ctypes.c_int


def hash_extend(orig_data_len, orig_hash, extend_data):
    if isinstance(orig_hash, str):
        orig_hash = binascii.unhexlify(orig_hash)

    total_len = math.ceil((orig_data_len + 1 + 8) / 64) * 64
    padding_len = total_len - orig_data_len - 1 - 8
    padding = b"\x80" + (b"\x00" * padding_len) + struct.pack(">Q", orig_data_len*8)

    ctx = sha2_ctx()
    sodium.crypto_hash_sha256_init(ctx)
    for i in range(8):
        ctx.state[i] = struct.unpack(">I", orig_hash[i*4:(i+1)*4])[0]
    ctx.count = total_len * 8
    out_hash = ctypes.create_string_buffer(32)
    sodium.crypto_hash_sha256_update(ctx, extend_data, len(extend_data))
    sodium.crypto_hash_sha256_final(ctx, out_hash)

    return padding + extend_data, bytes(out_hash)
