path: root/rand.py

#!/usr/bin/env python3

import sys
from typing import Optional

from cffi import FFI

class RandError(RuntimeError):
    pass

ffi = FFI()
ffi.cdef('''
    void RAND_seed(const void *buf, int num);

    void RAND_add(const void *buf, int num, double entropy);

    int  RAND_status(void);

    int  RAND_event(unsigned int iMsg, void* wParam, void* lParam);
    void RAND_screen(void);

    int RAND_bytes(unsigned char *buf, int num);

    int RAND_pseudo_bytes(unsigned char *buf, int num);

    const char *ERR_reason_error_string(unsigned long);

    unsigned long ERR_get_error(void);
''')

lib = ffi.dlopen('crypto')

def rand_add(buf: bytes, num: int, entropy: float) -> None:
    assert isinstance(entropy, float)
    lib.RAND_add(buf, num, entropy)

def rand_seed(buf: bytes, num: int) -> None:
    lib.RAND_seed(buf, num)

def rand_status() -> int:
    return lib.RAND_status()

def rand_event(iMsg: int, wParam: int, lParam: int) -> Optional[int]:
    if sys.platform == 'win32':
        return lib.RAND_event(iMsg, wParam, lParam)
    else:
        return None

def rand_screen() -> None:
    if sys.platform == 'win32':
        lib.RAND_screen()

def rand_bytes(n: int = 8) -> bytes:
    buf = ffi.new('unsigned char buf[{:d}]'.format(n))
    ret = lib.RAND_bytes(buf, n)
    if ret == 1:
        return bytes(buf)
    elif ret == 0:
        RandError("Not enough randomness.")
    elif ret == -1:
        RandError("Not supported by the current RAND method.")
    else:
        err_msg = lib.ERR_reason_error_string(lib.ERR_get_error())
        if err_msg is not None:
            RandError(err_msg)
        else:
            RandError("Unknown error in function RAND_bytes")

def rand_pseudo_bytes(n: int) -> bytes:
    blob = ffi.new('unsigned char buf[{:d}]'.format(n))
    ret = lib.RAND_pseudo_bytes(blob, n)

    if ret == -1:
        RandError("Not supported by the current RAND method.")
    elif ret in [0, 1]:
        return (bytes(blob), ret)
    else:
        RandError("Unknown error in RAND_pseudo_bytes")