cli_common.py - mozsearch

Enable keyboard shortcuts

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()