# procutil.py - utility for managing processes and executable environment

#

#  Copyright 2005 K. Thananchayan <thananck@yahoo.com>

#  Copyright 2005-2007 Olivia Mackall <olivia@selenic.com>

#  Copyright 2006 Vadim Gelfer <vadim.gelfer@gmail.com>

#

# This software may be used and distributed according to the terms of the

# GNU General Public License version 2 or any later version.

from __future__ import annotations

import contextlib

import errno

import io

import os

import signal

import subprocess

import sys

import threading

import time

from typing import (

    BinaryIO,

)

from ..i18n import _

from .. import (

    encoding,

    error,

    policy,

    pycompat,

    typelib,

)

# Import like this to keep import-checker happy

from ..utils import resourceutil

osutil = policy.importmod('osutil')

if pycompat.iswindows:

    from .. import windows as platform

else:

    from .. import posix as platform

def isatty(fp):

    try:

        return fp.isatty()

    except AttributeError:

        return False

class BadFile(io.RawIOBase):

    """Dummy file object to simulate closed stdio behavior"""

    def readinto(self, b):

        raise OSError(errno.EBADF, 'Bad file descriptor')

    def write(self, b):

        raise OSError(errno.EBADF, 'Bad file descriptor')

class LineBufferedWrapper:

    def __init__(self, orig):

        self.orig = orig

    def __getattr__(self, attr):

        return getattr(self.orig, attr)

    def write(self, s):

        orig = self.orig

        res = orig.write(s)

        if s.endswith(b'\n'):

            orig.flush()

        return res

# pytype: disable=attribute-error

io.BufferedIOBase.register(LineBufferedWrapper)

# pytype: enable=attribute-error

def make_line_buffered(stream):

    # First, check if we need to wrap the stream.

    check_stream = stream

    while True:

        if isinstance(check_stream, WriteAllWrapper):

            check_stream = check_stream.orig

        elif pycompat.iswindows and isinstance(

            check_stream,

            # pytype: disable=module-attr

            platform.winstdout

            # pytype: enable=module-attr

        ):

            check_stream = check_stream.fp

        else:

            break

    if isinstance(check_stream, io.RawIOBase):

        # The stream is unbuffered, we don't need to emulate line buffering.

        return stream

    elif isinstance(check_stream, io.BufferedIOBase):

        # The stream supports some kind of buffering. We can't assume that

        # lines are flushed. Fall back to wrapping the stream.

        pass

    else:

        raise NotImplementedError(

            "can't determine whether stream is buffered or not"

        )

    if isinstance(stream, LineBufferedWrapper):

        return stream

    return LineBufferedWrapper(stream)

def unwrap_line_buffered(stream):

    if isinstance(stream, LineBufferedWrapper):

        assert not isinstance(stream.orig, LineBufferedWrapper)

        return stream.orig

    return stream

class WriteAllWrapper(typelib.BinaryIO_Proxy):

    def __init__(self, orig: BinaryIO):

        self.orig = orig

    def __getattr__(self, attr):

        return getattr(self.orig, attr)

    def write(self, s):

        write1 = self.orig.write

        m = memoryview(s)

        total_to_write = len(s)

        total_written = 0

        while total_written < total_to_write:

            c = write1(m[total_written:])

            if c:

                total_written += c

        return total_written

# pytype: disable=attribute-error

io.IOBase.register(WriteAllWrapper)

# pytype: enable=attribute-error

def _make_write_all(stream):

    if isinstance(stream, WriteAllWrapper):

        return stream

    if isinstance(stream, io.BufferedIOBase):

        # The io.BufferedIOBase.write() contract guarantees that all data is

        # written.

        return stream

    # In general, the write() method of streams is free to write only part of

    # the data.

    return WriteAllWrapper(stream)

# Python 3 implements its own I/O streams. Unlike stdio of C library,

# sys.stdin/stdout/stderr may be None if underlying fd is closed.

# TODO: .buffer might not exist if std streams were replaced; we'll need

# a silly wrapper to make a bytes stream backed by a unicode one.

if sys.stdin is None:

    stdin = BadFile()

else:

    stdin = sys.stdin.buffer

if sys.stdout is None:

    stdout = BadFile()

else:

    stdout = _make_write_all(sys.stdout.buffer)

if sys.stderr is None:

    stderr = BadFile()

else:

    stderr = _make_write_all(sys.stderr.buffer)

if pycompat.iswindows:

    # Work around Windows bugs.

    stdout = platform.winstdout(stdout)  # pytype: disable=module-attr

    stderr = platform.winstdout(stderr)  # pytype: disable=module-attr

if isatty(stdout) and pycompat.sysplatform != b'OpenVMS':

    # The standard library doesn't offer line-buffered binary streams.

    stdout = make_line_buffered(stdout)

findexe = platform.findexe

_gethgcmd = platform.gethgcmd

getuser = platform.getuser

getpid = os.getpid

hidewindow = platform.hidewindow

readpipe = platform.readpipe

setbinary = platform.setbinary

setsignalhandler = platform.setsignalhandler

shellquote = platform.shellquote

shellsplit = platform.shellsplit

spawndetached = platform.spawndetached

sshargs = platform.sshargs

testpid = platform.testpid

try:

    setprocname = osutil.setprocname

except AttributeError:

    pass

try:

    unblocksignal = osutil.unblocksignal

except AttributeError:

    pass

closefds = pycompat.isposix and pycompat.sysplatform != b'OpenVMS'

def explainexit(code):

    """return a message describing a subprocess status

    (codes from kill are negative - not os.system/wait encoding)"""

    if code >= 0:

        return _(b"exited with status %d") % code

    return _(b"killed by signal %d") % -code

class _pfile:

    """File-like wrapper for a stream opened by subprocess.Popen()"""

    def __init__(self, proc, fp):

        self._proc = proc

        self._fp = fp

    def close(self):

        # unlike os.popen(), this returns an integer in subprocess coding

        self._fp.close()

        return self._proc.wait()

    def __iter__(self):

        return iter(self._fp)

    def __getattr__(self, attr):

        return getattr(self._fp, attr)

    def __enter__(self):

        return self

    def __exit__(self, exc_type, exc_value, exc_tb):

        self.close()

def popen(cmd, mode=b'rb', bufsize=-1):

    if mode == b'rb':

        return _popenreader(cmd, bufsize)

    elif mode == b'wb':

        return _popenwriter(cmd, bufsize)

    raise error.ProgrammingError(b'unsupported mode: %r' % mode)

def _popenreader(cmd, bufsize):

    p = subprocess.Popen(

        tonativestr(cmd),

        shell=True,

        bufsize=bufsize,

        close_fds=closefds,

        stdout=subprocess.PIPE,

    )

    return _pfile(p, p.stdout)

def _popenwriter(cmd, bufsize):

    p = subprocess.Popen(

        tonativestr(cmd),

        shell=True,

        bufsize=bufsize,

        close_fds=closefds,

        stdin=subprocess.PIPE,

    )

    return _pfile(p, p.stdin)

def popen2(cmd, env=None):

    # Setting bufsize to -1 lets the system decide the buffer size.

    # The default for bufsize is 0, meaning unbuffered. This leads to

    # poor performance on Mac OS X: http://bugs.python.org/issue4194

    p = subprocess.Popen(

        tonativestr(cmd),

        shell=True,

        bufsize=-1,

        close_fds=closefds,

        stdin=subprocess.PIPE,

        stdout=subprocess.PIPE,

        env=tonativeenv(env),

    )

    return p.stdin, p.stdout

def popen3(cmd, env=None):

    stdin, stdout, stderr, p = popen4(cmd, env)

    return stdin, stdout, stderr

def popen4(cmd, env=None, bufsize=-1):

    p = subprocess.Popen(

        tonativestr(cmd),

        shell=True,

        bufsize=bufsize,

        close_fds=closefds,

        stdin=subprocess.PIPE,

        stdout=subprocess.PIPE,

        stderr=subprocess.PIPE,

        env=tonativeenv(env),

    )

    return p.stdin, p.stdout, p.stderr, p

def pipefilter(s, cmd):

    '''filter string S through command CMD, returning its output'''

    p = subprocess.Popen(

        tonativestr(cmd),

        shell=True,

        close_fds=closefds,

        stdin=subprocess.PIPE,

        stdout=subprocess.PIPE,

    )

    pout, perr = p.communicate(s)

    return pout

def tempfilter(s, cmd):

    """filter string S through a pair of temporary files with CMD.

    CMD is used as a template to create the real command to be run,

    with the strings INFILE and OUTFILE replaced by the real names of

    the temporary files generated."""

    inname, outname = None, None

    try:

        infd, inname = pycompat.mkstemp(prefix=b'hg-filter-in-')

        with os.fdopen(infd, 'wb') as fp:

            fp.write(s)

        outfd, outname = pycompat.mkstemp(prefix=b'hg-filter-out-')

        os.close(outfd)

        cmd = cmd.replace(b'INFILE', inname)

        cmd = cmd.replace(b'OUTFILE', outname)

        code = system(cmd)

        if code:

            raise error.Abort(

                _(b"command '%s' failed: %s") % (cmd, explainexit(code))

            )

        with open(outname, 'rb') as fp:

            return fp.read()

    finally:

        try:

            if inname:

                os.unlink(inname)

        except OSError:

            pass

        try:

            if outname:

                os.unlink(outname)

        except OSError:

            pass

_filtertable = {

    b'tempfile:': tempfilter,

    b'pipe:': pipefilter,

}

def filter(s, cmd):

    """filter a string through a command that transforms its input to its

    output"""

    for name, fn in _filtertable.items():

        if cmd.startswith(name):

            return fn(s, cmd[len(name) :].lstrip())

    return pipefilter(s, cmd)

_hgexecutable = None

def hgexecutable():

    """return location of the 'hg' executable.

    Defaults to $HG or 'hg' in the search path.

    """

    if _hgexecutable is None:

        hg = encoding.environ.get(b'HG', '')

        mainmod = sys.modules['__main__']

        if pycompat.sysplatform == b'OpenVMS' and hg[0:1] == '$':

            hg = 'mcr ' + hg[1:]

        if hg:

            _sethgexecutable(hg)

        elif resourceutil.mainfrozen():

            if getattr(sys, 'frozen', None) == 'macosx_app':

                # Env variable set by py2app

                _sethgexecutable(encoding.environ[b'EXECUTABLEPATH'])

            else:

                _sethgexecutable(pycompat.sysexecutable)

        elif (

            not pycompat.iswindows

            and os.path.basename(getattr(mainmod, '__file__', '')) == 'hg'

        ):

            _sethgexecutable(pycompat.fsencode(mainmod.__file__))

        else:

            _sethgexecutable(

                findexe(b'hg') or os.path.basename(pycompat.sysargv[0])

            )

    return _hgexecutable

def _sethgexecutable(path):

    """set location of the 'hg' executable"""

    global _hgexecutable

    _hgexecutable = path

def _testfileno(f, stdf):

    fileno = getattr(f, 'fileno', None)

    try:

        return fileno and fileno() == stdf.fileno()

    except io.UnsupportedOperation:

        return False  # fileno() raised UnsupportedOperation

def isstdin(f):

    return _testfileno(f, sys.__stdin__)

def isstdout(f):

    return _testfileno(f, sys.__stdout__)

def protectstdio(uin, uout):

    """Duplicate streams and redirect original if (uin, uout) are stdio

    If uin is stdin, it's redirected to /dev/null. If uout is stdout, it's

    redirected to stderr so the output is still readable.

    Returns (fin, fout) which point to the original (uin, uout) fds, but

    may be copy of (uin, uout). The returned streams can be considered

    "owned" in that print(), exec(), etc. never reach to them.

    """

    uout.flush()

    fin, fout = uin, uout

    if _testfileno(uin, stdin):

        newfd = os.dup(uin.fileno())

        nullfd = os.open(os.devnull, os.O_RDONLY)

        os.dup2(nullfd, uin.fileno())

        os.close(nullfd)

        fin = os.fdopen(newfd, 'rb')

    if _testfileno(uout, stdout):

        newfd = os.dup(uout.fileno())

        os.dup2(stderr.fileno(), uout.fileno())

        fout = os.fdopen(newfd, 'wb')

    return fin, fout

def restorestdio(uin, uout, fin, fout):

    """Restore (uin, uout) streams from possibly duplicated (fin, fout)"""

    uout.flush()

    for f, uif in [(fin, uin), (fout, uout)]:

        if f is not uif:

            os.dup2(f.fileno(), uif.fileno())

            f.close()

def shellenviron(environ=None):

    """return environ with optional override, useful for shelling out"""

    def py2shell(val):

        """convert python object into string that is useful to shell"""

        if val is None or val is False:

            return b'0'

        if val is True:

            return b'1'

        return pycompat.bytestr(val)

    env = dict(encoding.environ)

    if environ:

        env.update((k, py2shell(v)) for k, v in environ.items())

    env[b'HG'] = hgexecutable()

    return env