path: root/src/tests/cli_common.py

import sys
import random
import string
import logging
import os
import re
import shutil
from subprocess import Popen, PIPE

RNP_ROOT = None
WORKDIR = ''
CONSOLE_ENCODING = 'UTF-8'

class CLIError(Exception):
    def __init__(self, message, log = None):
        super(CLIError, self).__init__(message)
        self.message = message
        self.log = log
        logging.info(self.message)
        logging.debug(self.log.strip())

    def __str__(self):
        return self.message + '\n' + self.log

def set_workdir(dir):
    global WORKDIR
    WORKDIR = dir

def is_windows():
    return sys.platform.startswith('win') or sys.platform.startswith('msys')

def path_for_gpg(path):
    # GPG built for mingw/msys doesn't work with Windows paths
    if re.match(r'^[a-z]:[\\\/].*', path.lower()):
        path = '/' + path[0] + '/' + path[3:].replace('\\', '/')
    return path

def raise_err(msg, log = None):
    raise CLIError(msg, log)

def size_to_readable(num, suffix = 'B'):
    for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:
        if abs(num) < 1024.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s" % (num, 'Yi', suffix)

def list_upto(lst, count):
    return (list(lst)*(count//len(lst)+1))[:count]

def pswd_pipe(password):
    pr, pw = os.pipe()
    with os.fdopen(pw, 'w') as fw:
        fw.write(password)
        fw.write('\n')
        fw.write(password)
    os.set_inheritable(pr, True)

    if not is_windows():
        return pr
    # On Windows pipe is not inheritable so dup() is needed
    prd = os.dup(pr)
    os.close(pr)
    return prd

def random_text(path, size):
    # Generate random text, with 50% probability good-compressible
    if random.randint(0, 10) < 5:
        st = ''.join(random.choice(string.ascii_letters + string.digits + " \t\r\n-,.")
                     for _ in range(size))
    else:
        st = ''.join(random.choice("abcdef0123456789 \t\r\n-,.") for _ in range(size))
    with open(path, 'w+') as f:
        f.write(st)

def file_text(path, encoding = CONSOLE_ENCODING):
    with open(path, 'rb') as f:
        return f.read().decode(encoding).replace('\r\r', '\r')

def find_utility(name, exitifnone = True):
    path = shutil.which(name)
    if not path and exitifnone:
        logging.error('Cannot find utility {}. Exiting.'.format(name))
        sys.exit(1)

    return path

def rnp_file_path(relpath, check = True):
    global RNP_ROOT
    if not RNP_ROOT:
        pypath = os.path.dirname(__file__)
        RNP_ROOT = os.path.realpath(os.path.join(pypath, '../..'))

    fpath = os.path.realpath(os.path.join(RNP_ROOT, relpath))

    if check and not os.path.isfile(fpath):
        raise NameError('rnp: file ' + relpath + ' not found')

    return fpath

def run_proc_windows(proc, params, stdin=None):
    exe = os.path.basename(proc)
    # test special quote cases 
    params = list(map(lambda st: st.replace('"', '\\"'), params))
    # We need to escape empty parameters/ones with spaces with quotes
    params = tuple(map(lambda st: st if (st and not any(x in st for x in [' ','\r','\t'])) else '"%s"' % st, [exe] + params))
    logging.debug((proc + ' ' + ' '.join(params)).strip())
    logging.debug('Working directory: ' + os.getcwd())
    sys.stdout.flush()

    stdin_path = os.path.join(WORKDIR, 'stdin.txt')
    stdout_path = os.path.join(WORKDIR, 'stdout.txt')
    stderr_path = os.path.join(WORKDIR, 'stderr.txt')
    pass_path = os.path.join(WORKDIR, 'pass.txt')
    passfd = 0
    passfo = None
    try:
        idx = params.index('--pass-fd')
        if idx < len(params):
            passfd = int(params[idx+1])
            passfo = os.fdopen(passfd, 'r', closefd=False)
    except (ValueError, OSError):
        # Ignore if pass-fd is invalid/could not be opened
        pass
    # We may use pipes here (ensuring we use dup to inherit handles), but those have limited buffer
    # so we'll need to poll process
    if stdin:
        with open(stdin_path, "wb+") as stdinf:
            stdinf.write(stdin.encode() if isinstance(stdin, str) else stdin)
        stdin_fl = os.open(stdin_path, os.O_RDONLY | os.O_BINARY)
        stdin_no = sys.stdin.fileno()
        stdin_cp = os.dup(stdin_no)
    else:
        stdin_fl = None
        stdin_no = -1
        stdin_cp = None

    stdout_fl = os.open(stdout_path, os.O_CREAT | os.O_RDWR | os.O_BINARY)
    stdout_no = sys.stdout.fileno()
    stdout_cp = os.dup(stdout_no)
    stderr_fl = os.open(stderr_path, os.O_CREAT | os.O_RDWR | os.O_BINARY)
    stderr_no = sys.stderr.fileno()
    stderr_cp = os.dup(stderr_no)
    if passfo:
        with open(pass_path, "w+") as passf:
            passf.write(passfo.read())
        pass_fl = os.open(pass_path, os.O_RDONLY | os.O_BINARY)
        pass_cp = os.dup(passfd)

    retcode = -1
    try:
        os.dup2(stdout_fl, stdout_no)
        os.close(stdout_fl)
        os.dup2(stderr_fl, stderr_no)
        os.close(stderr_fl)
        if stdin:
            os.dup2(stdin_fl, stdin_no)
            os.close(stdin_fl)
        if passfo:
            os.dup2(pass_fl, passfd)
            os.close(pass_fl)
        retcode = os.spawnv(os.P_WAIT, proc, params)
    finally:
        os.dup2(stdout_cp, stdout_no)
        os.close(stdout_cp)
        os.dup2(stderr_cp, stderr_no)
        os.close(stderr_cp)
        if stdin:
            os.dup2(stdin_cp, stdin_no)
            os.close(stdin_cp)
        if passfo:
            os.dup2(pass_cp, passfd)
            os.close(pass_cp)
            passfo.close()
    out = file_text(stdout_path).replace('\r\n', '\n')
    err = file_text(stderr_path).replace('\r\n', '\n')
    os.unlink(stdout_path)
    os.unlink(stderr_path)
    if stdin: 
        os.unlink(stdin_path)
    if passfo: 
        os.unlink(pass_path)
    logging.debug(err.strip())
    logging.debug(out.strip())
    return (retcode, out, err)

if sys.version_info >= (3,):
    def decode_string_escape(s):
        bts = bytes(s, 'utf-8')
        result = u''
        candidate = bytearray()
        utf = bytearray()
        for b in bts:
            if b > 0x7F:
                if len(candidate) > 0:
                    result += candidate.decode('unicode-escape')
                    candidate.clear()
                utf.append(b)
            else:
                if len(utf) > 0:
                    result += utf.decode('utf-8')
                    utf.clear()
                candidate.append(b)
        if len(candidate) > 0:
            result += candidate.decode('unicode-escape')
        if len(utf) > 0:
            result += utf.decode('utf-8')
        return result
    def _decode(s):
        return s
else: # Python 2
    def decode_string_escape(s):
        return s.encode(CONSOLE_ENCODING).decode('decode_string_escape')
    def _decode(x):
        return x.decode(CONSOLE_ENCODING)

def run_proc(proc, params, stdin=None):
    # On Windows we need to use spawnv() for handle inheritance in pswd_pipe()
    if is_windows():
        return run_proc_windows(proc, params, stdin)
    paramline = u' '.join(map(_decode, params))
    logging.debug((proc + ' ' + paramline).strip())
    param_bytes = list(map(lambda x: x.encode(CONSOLE_ENCODING), params))
    process = Popen([proc] + param_bytes, stdout=PIPE, stderr=PIPE,
                    stdin=PIPE if stdin else None, close_fds=False,
                    universal_newlines=True)
    output, errout = process.communicate(stdin)
    retcode = process.poll()
    logging.debug(errout.strip())
    logging.debug(output.strip())

    return (retcode, output, errout)

def run_proc_fast(proc, params):
    with open(os.devnull, 'w') as devnull:
        proc = Popen([proc] + params, stdout=devnull, stderr=devnull)
    return proc.wait()