#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Run processes of a Thermos task.

This module contains the Process class, used to manage the execution of the constituent processes of
a Thermos task. Each process is represented by a "coordinator" process, which fires off the actual
commandline in a subprocess of its own.

"""

import errno
import grp
import os
import pwd
import select
import signal
import subprocess
import sys
import time
from abc import abstractmethod
from copy import deepcopy

from twitter.common import log
from twitter.common.dirutil import lock_file, safe_delete, safe_mkdir, safe_open
from twitter.common.lang import Interface
from twitter.common.quantity import Amount, Data, Time
from twitter.common.recordio import ThriftRecordReader, ThriftRecordWriter

from apache.thermos.common.process_util import setup_child_subreaping, wrap_with_mesos_containerizer

from gen.apache.aurora.api.constants import TASK_FILESYSTEM_MOUNT_POINT
from gen.apache.thermos.ttypes import ProcessState, ProcessStatus, RunnerCkpt


class Platform(Interface):
  """Abstract representation of a platform encapsulating system-level functions"""

  @abstractmethod
  def clock(self):
    pass

  @abstractmethod
  def fork(self):
    pass

  @abstractmethod
  def getpid(self):
    pass


class LoggerDestination(object):
  FILE = 'file'
  CONSOLE = 'console'
  BOTH = 'both'
  NONE = 'none'

  _ALL_DESTINATIONS = [FILE, CONSOLE, BOTH, NONE]

  @staticmethod
  def is_valid(destination):
    return destination in LoggerDestination._ALL_DESTINATIONS


class LoggerMode(object):
  STANDARD = 'standard'
  ROTATE = 'rotate'

  _ALL_MODES = [STANDARD, ROTATE]

  @staticmethod
  def is_valid(mode):
    return mode in LoggerMode._ALL_MODES


class ProcessBase(object):
  """
    Encapsulate a running process for a task.
  """
  class Error(Exception): pass
  class UnknownUserError(Error): pass
  class CheckpointError(Error): pass
  class UnspecifiedSandbox(Error): pass
  class PermissionError(Error): pass

  CONTROL_WAIT_CHECK_INTERVAL = Amount(100, Time.MILLISECONDS)
  MAXIMUM_CONTROL_WAIT = Amount(1, Time.MINUTES)

  def __init__(self, name, cmdline, sequence, pathspec, sandbox_dir, user=None, platform=None,
               logger_destination=LoggerDestination.FILE, logger_mode=LoggerMode.STANDARD,
               rotate_log_size=None, rotate_log_backups=None):
    """
      required:
        name        = name of the process
        cmdline     = cmdline of the process
        sequence    = the next available sequence number for state updates
        pathspec    = TaskPath object for synthesizing path names
        sandbox_dir = the sandbox in which to run the process
        platform    = Platform providing fork, clock, getpid

      optional:
        user               = the user to run as (if unspecified, will default to current user.)
                             if specified to a user that is not the current user, you must have root
                             access
        logger_destination = The destination for logs output.
        logger_mode        = The type of logger to use for the process.
        rotate_log_size    = The maximum size of the rotated stdout/stderr logs.
        rotate_log_backups = The maximum number of rotated stdout/stderr log backups.
    """
    self._name = name
    self._cmdline = cmdline
    self._pathspec = pathspec
    self._seq = sequence
    self._sandbox = sandbox_dir
    if self._sandbox:
      safe_mkdir(self._sandbox)
    self._pid = None
    self._fork_time = None
    self._user = user
    self._ckpt = None
    self._ckpt_head = -1
    if platform is None:
      raise ValueError("Platform must be specified")
    self._platform = platform
    self._logger_destination = logger_destination
    self._logger_mode = logger_mode
    self._rotate_log_size = rotate_log_size
    self._rotate_log_backups = rotate_log_backups

    if not LoggerDestination.is_valid(self._logger_destination):
      raise ValueError("Logger destination %s is invalid." % self._logger_destination)

    if not LoggerMode.is_valid(self._logger_mode):
      raise ValueError("Logger mode %s is invalid." % self._logger_mode)

    if self._logger_mode == LoggerMode.ROTATE:
      if self._rotate_log_size.as_(Data.BYTES) <= 0:
        raise ValueError('Log size cannot be less than one byte.')
      if self._rotate_log_backups <= 0:
        raise ValueError('Log backups cannot be less than one.')

  def _log(self, msg, exc_info=None):
    log.debug('[process:%5s=%s]: %s' % (self._pid, self.name(), msg),
            exc_info=exc_info)

  def _getpwuid(self):
    """Returns a tuple of the user (i.e. --user) and current user."""
    uid = os.getuid()
    try:
      current_user = pwd.getpwuid(uid)
    except KeyError:
      raise self.UnknownUserError('Unknown uid %s!' % uid)
    try:
      user = pwd.getpwnam(self._user) if self._user is not None else current_user
    except KeyError:
      raise self.UnknownUserError('Unable to get pwent information!')
    return user, current_user

  def _ckpt_write(self, msg):
    self._init_ckpt_if_necessary()
    self._log("child state transition [%s] <= %s" % (self.ckpt_file(), msg))
    self._ckpt.write(msg)

  def _write_process_update(self, **kw):
    """Write a process update to the coordinator's checkpoint stream."""
    process_status = ProcessStatus(**kw)
    process_status.seq = self._seq
    process_status.process = self.name()
    self._ckpt_write(RunnerCkpt(process_status=process_status))
    self._seq += 1

  def _write_initial_update(self):
    self._write_process_update(state=ProcessState.FORKED,
                               fork_time=self._fork_time,
                               coordinator_pid=self._pid)

  def cmdline(self):
    return self._cmdline

  def name(self):
    return self._name

  def pid(self):
    """pid of the coordinator"""
    return self._pid

  def rebind(self, pid, fork_time):
    """rebind Process to an existing coordinator pid without forking"""
    self._pid = pid
    self._fork_time = fork_time

  def ckpt_file(self):
    return self._pathspec.getpath('process_checkpoint')

  def process_logdir(self):
    return self._pathspec.getpath('process_logdir')

  def _setup_ckpt(self):
    """Set up the checkpoint: must be run on the parent."""
    self._log('initializing checkpoint file: %s' % self.ckpt_file())
    ckpt_fp = lock_file(self.ckpt_file(), "a+")
    if ckpt_fp in (None, False):
      raise self.CheckpointError('Could not acquire checkpoint permission or lock for %s!' %
        self.ckpt_file())
    self._ckpt_head = os.path.getsize(self.ckpt_file())
    ckpt_fp.seek(self._ckpt_head)
    self._ckpt = ThriftRecordWriter(ckpt_fp)
    self._ckpt.set_sync(True)

  def _init_ckpt_if_necessary(self):
    if self._ckpt is None:
      self._setup_ckpt()

  def _wait_for_control(self):
    """Wait for control of the checkpoint stream: must be run in the child."""
    total_wait_time = Amount(0, Time.SECONDS)

    with open(self.ckpt_file(), 'r') as fp:
      fp.seek(self._ckpt_head)
      rr = ThriftRecordReader(fp, RunnerCkpt)
      while total_wait_time < self.MAXIMUM_CONTROL_WAIT:
        ckpt_tail = os.path.getsize(self.ckpt_file())
        if ckpt_tail == self._ckpt_head:
          self._platform.clock().sleep(self.CONTROL_WAIT_CHECK_INTERVAL.as_(Time.SECONDS))
          total_wait_time += self.CONTROL_WAIT_CHECK_INTERVAL
          continue
        checkpoint = rr.try_read()
        if checkpoint:
          if not checkpoint.process_status:
            raise self.CheckpointError('No process status in checkpoint!')
          if (checkpoint.process_status.process != self.name() or
              checkpoint.process_status.state != ProcessState.FORKED or
              checkpoint.process_status.fork_time != self._fork_time or
              checkpoint.process_status.coordinator_pid != self._pid):
            self._log('Losing control of the checkpoint stream:')
            self._log('   fork_time [%s] vs self._fork_time [%s]' % (
                checkpoint.process_status.fork_time, self._fork_time))
            self._log('   coordinator_pid [%s] vs self._pid [%s]' % (
                checkpoint.process_status.coordinator_pid, self._pid))
            raise self.CheckpointError('Lost control of the checkpoint stream!')
          self._log('Taking control of the checkpoint stream at record: %s' %
            checkpoint.process_status)
          self._seq = checkpoint.process_status.seq + 1
          return True
    raise self.CheckpointError('Timed out waiting for checkpoint stream!')

  def _prepare_fork(self):
    user, current_user = self._getpwuid()
    if self._user:
      if user != current_user and os.geteuid() != 0:
        raise self.PermissionError('Must be root to run processes as other users!')
    self._fork_time = self._platform.clock().time()
    self._setup_ckpt()
    # Since the forked process is responsible for creating log files, it needs to own the log dir.
    safe_mkdir(self.process_logdir())
    os.chown(self.process_logdir(), user.pw_uid, user.pw_gid)

  def _finalize_fork(self):
    self._write_initial_update()
    self._ckpt.close()
    self._ckpt = None

  def start(self):
    """
      This is the main call point from the runner, and forks a co-ordinator process to run the
      target process (i.e. self.cmdline())

      The parent returns immediately and populates information about the pid of the co-ordinator.
      The child (co-ordinator) will launch the target process in a subprocess.
    """
    self._prepare_fork()  # calls _setup_ckpt which can raise CheckpointError
                          # calls _getpwuid which can raise:
                          #    UnknownUserError
                          #    PermissionError
    self._pid = self._platform.fork()
    if self._pid == 0:
      self._pid = self._platform.getpid()
      self._wait_for_control()  # can raise CheckpointError
      try:
        self.execute()
      except Exception as e:
        self._log('Error trying to execute %s: %s' % (self._name, e))
        raise e
      finally:
        self._ckpt.close()
        self.finish()
    else:
      self._finalize_fork()  # can raise CheckpointError

  def execute(self):
    raise NotImplementedError

  def finish(self):
    pass


class RealPlatform(Platform):
  IGNORE_SIGNALS = (signal.SIGINT,)

  def __init__(self, fork=os.fork):
    self._fork = fork

  def fork(self):
    # Before we fork, ensure we become the parent of any processes that escape
    # the cordinator.
    setup_child_subreaping()
    pid = self._fork()
    if pid == 0:
      self._sanitize()
    return pid

  def _sanitize(self):
    for sig in self.IGNORE_SIGNALS:
      signal.signal(sig, signal.SIG_IGN)

  def getpid(self):
    return os.getpid()

  def clock(self):
    return time


class Process(ProcessBase):
  """
    Encapsulate a running process for a task.
  """
  RCFILE = '.thermos_profile'
  FD_CLOEXEC = True

  def __init__(self, *args, **kw):
    """
      See ProcessBase.__init__

      Takes additional arguments:
        fork: the fork function to use [default: os.fork]
        chroot: whether or not to chroot into the sandbox [default: False]
        preserve_env: whether or not to preserve env variables for the task [default: False]
        mesos_containerizer_path: The path to the mesos-containerizer binary to be used for task
                                  filesystem isolation.
        container_sandbox: If running in an isolated filesystem, the path within that filesystem
                           where the sandbox is mounted.
    """
    fork = kw.pop('fork', os.fork)
    self._use_chroot = bool(kw.pop('chroot', False))
    self._rc = None
    self._preserve_env = bool(kw.pop('preserve_env', False))
    self._mesos_containerizer_path = kw.pop('mesos_containerizer_path', None)
    self._container_sandbox = kw.pop('container_sandbox', None)

    if self._mesos_containerizer_path is not None and self._container_sandbox is None:
      raise self.UnspecifiedSandbox('If using mesos-containerizer, container_sandbox must be set.')

    kw['platform'] = RealPlatform(fork=fork)
    ProcessBase.__init__(self, *args, **kw)
    if self._use_chroot and self._sandbox is None:
      raise self.UnspecifiedSandbox('If using chroot, must specify sandbox!')

  def _chroot(self):
    """chdir and chroot to the sandbox directory."""
    os.chdir(self._sandbox)
    os.chroot(self._sandbox)

  def _setuid(self):
    """Drop privileges to the user supplied in Process creation (if necessary.)"""
    user, current_user = self._getpwuid()
    if user.pw_uid == current_user.pw_uid:
      return

    uid, gid = user.pw_uid, user.pw_gid
    username = user.pw_name
    group_ids = [group.gr_gid for group in grp.getgrall() if username in group.gr_mem]
    os.setgroups(group_ids)
    os.setgid(gid)
    os.setuid(uid)

  def wrapped_cmdline(self, cwd):
    cmdline = self.cmdline()

    # If mesos-containerizer is not set, we only need to wrap the cmdline in a bash invocation.
    if self._mesos_containerizer_path is None:
      return ['/bin/bash', '-c', cmdline]
    else:
      return wrap_with_mesos_containerizer(cmdline, self._user, cwd, self._mesos_containerizer_path)

  def execute(self):
    """Perform final initialization and launch target process commandline in a subprocess."""

    user, _ = self._getpwuid()
    username, homedir = user.pw_name, user.pw_dir

    # TODO(wickman) reconsider setsid now that we're invoking in a subshell
    os.setsid()
    if self._use_chroot:
      self._chroot()

    # If the mesos containerizer path is set, then this process will be launched from within an
    # isolated filesystem image by the mesos-containerizer executable. This executable needs to be
    # run as root so that it can properly set up the filesystem as such we'll skip calling setuid at
    # this point. We'll instead setuid after the process has been forked (mesos-containerizer itself
    # ensures the forked process is run as the correct user).
    taskfs_isolated = self._mesos_containerizer_path is not None
    if not taskfs_isolated:
      self._setuid()

    # start process
    start_time = self._platform.clock().time()

    if not self._sandbox:
      cwd = subprocess_cwd = sandbox = os.getcwd()
    else:
      if self._use_chroot:
        cwd = subprocess_cwd = sandbox = '/'
      elif taskfs_isolated:
        cwd = homedir = sandbox = self._container_sandbox
        subprocess_cwd = self._sandbox
      else:
        cwd = subprocess_cwd = homedir = sandbox = self._sandbox

    thermos_profile = os.path.join(sandbox, self.RCFILE)

    if self._preserve_env:
      env = deepcopy(os.environ)
    else:
      env = {}

    env.update({
      'HOME': homedir,
      'LOGNAME': username,
      'USER': username,
      'PATH': os.environ['PATH']
    })

    wrapped_cmdline = self.wrapped_cmdline(cwd)
    log.debug('Wrapped cmdline: %s' % wrapped_cmdline)

    real_thermos_profile_path = os.path.join(
        os.environ['MESOS_DIRECTORY'],
        TASK_FILESYSTEM_MOUNT_POINT,
        thermos_profile.lstrip('/')) if taskfs_isolated else thermos_profile

    if os.path.exists(real_thermos_profile_path):
      env.update(BASH_ENV=thermos_profile)

    log.debug('ENV is: %s' % env)
    subprocess_args = {
      'args': wrapped_cmdline,
      'close_fds': self.FD_CLOEXEC,
      'cwd': subprocess_cwd,
      'env': env,
      'pathspec': self._pathspec
    }

    log_destination_resolver = LogDestinationResolver(
        self._pathspec,
        destination=self._logger_destination,
        mode=self._logger_mode,
        rotate_log_size=self._rotate_log_size,
        rotate_log_backups=self._rotate_log_backups)
    stdout, stderr, handlers_are_files = log_destination_resolver.get_handlers()
    if handlers_are_files:
      executor = SubprocessExecutor(stdout=stdout, stderr=stderr, **subprocess_args)
    else:
      executor = PipedSubprocessExecutor(stdout=stdout, stderr=stderr, **subprocess_args)

    pid = executor.start()

    # Now that we've forked the process, if the task's filesystem is isolated it's now safe to
    # setuid.
    if taskfs_isolated:
      self._setuid()

    self._write_process_update(state=ProcessState.RUNNING, pid=pid, start_time=start_time)

    rc = executor.wait()

    # indicate that we have finished/failed
    if rc < 0:
      state = ProcessState.KILLED
    elif rc == 0:
      state = ProcessState.SUCCESS
    else:
      state = ProcessState.FAILED

    self._write_process_update(state=state, return_code=rc, stop_time=self._platform.clock().time())
    self._rc = rc

  def finish(self):
    self._log('Coordinator exiting.')
    sys.exit(0)


class SubprocessExecutorBase(object):
  """
  Encapsulate execution of a subprocess.
  """

  def __init__(self, args, close_fds, cwd, env, pathspec):
    """
      required:
        args        = The arguments to pass to the subprocess.
        close_fds   = Close file descriptors argument to Popen.
        cwd         = The current working directory.
        env         = Environment variables to be passed to the subprocess.
        pathspec    = TaskPath object for synthesizing path names.
    """
    self._args = args
    self._close_fds = close_fds
    self._cwd = cwd
    self._env = env
    self._pathspec = pathspec
    self._popen = None

  def _start_subprocess(self, stderr, stdout):
    return subprocess.Popen(self._args,
                            stderr=stderr,
                            stdout=stdout,
                            close_fds=self._close_fds,
                            cwd=self._cwd,
                            env=self._env)

  def start(self):
    """Start the subprocess and immediately return the resulting pid."""
    raise NotImplementedError()

  def wait(self):
    """Wait for the subprocess to finish executing and return the return code."""
    raise NotImplementedError()


class SubprocessExecutor(SubprocessExecutorBase):
  """
  Basic implementation of a SubprocessExecutor that writes stderr/stdout to specified output files.
  """

  def __init__(self, args, close_fds, cwd, env, pathspec, stdout=None, stderr=None):
    """
    See SubprocessExecutorBase.__init__

    Takes additional arguments:
      stdout = Destination handler for stdout output. Default is /dev/null.
      stderr = Destination handler for stderr output. Default is /dev/null.
    """
    super(SubprocessExecutor, self).__init__(args, close_fds, cwd, env, pathspec)
    self._stderr = stderr
    self._stdout = stdout

  def start(self):
    self._popen = self._start_subprocess(self._stderr, self._stdout)
    return self._popen.pid

  def wait(self):
    return self._popen.wait()


class PipedSubprocessExecutor(SubprocessExecutorBase):
  """
  Implementation of SubprocessExecutorBase that uses pipes to poll the pipes to output streams and
  copies them to the specified destinations.
  """

  READ_BUFFER_SIZE = 2 ** 16

  def __init__(self, args, close_fds, cwd, env, pathspec, stdout=None, stderr=None):
    """
    See SubprocessExecutorBase.__init__

    Takes additional arguments:
      stdout = Destination handler for stdout output. Default is /dev/null.
      stderr = Destination handler for stderr output. Default is /dev/null.
    """
    super(PipedSubprocessExecutor, self).__init__(args, close_fds, cwd, env, pathspec)
    self._stderr = stderr
    self._stdout = stdout

  def start(self):
    self._popen = self._start_subprocess(subprocess.PIPE, subprocess.PIPE)
    return self._popen.pid

  def wait(self):
    stdout = self._popen.stdout.fileno()
    stderr = self._popen.stderr.fileno()
    pipes = {
      stderr: self._stderr,
      stdout: self._stdout
    }

    rc = None
    # Read until there is a return code AND both of the pipes have reached EOF.
    while rc is None or pipes:
      rc = self._popen.poll()

      read_results, _, _ = select.select(pipes.keys(), [], [], 1)
      for fd in read_results:
        handler = pipes[fd]
        buf = os.read(fd, self.READ_BUFFER_SIZE)

        if len(buf) == 0:
          del pipes[fd]
        else:
          handler.write(buf)

    return rc


class LogDestinationResolver(object):
  """
  Resolves correct stdout/stderr destinations based on process configuration.
  """

  STDOUT = 'stdout'
  STDERR = 'stderr'

  def __init__(self, pathspec, destination=LoggerDestination.FILE, mode=LoggerMode.STANDARD,
               rotate_log_size=None, rotate_log_backups=None):
    """
    pathspec           = TaskPath object for synthesizing path names.
    destination        = Log destination.
    logger_mode        = The type of logger to use for the process.
    rotate_log_size    = The maximum size of the rotated stdout/stderr logs.
    rotate_log_backups = The maximum number of rotated stdout/stderr log backups.
    """
    self._pathspec = pathspec
    self._destination = destination
    self._mode = mode
    self._rotate_log_size = rotate_log_size
    self._rotate_log_backups = rotate_log_backups

    if not LoggerDestination.is_valid(self._destination):
      raise ValueError("Logger destination %s is invalid." % self._destination)

    if not LoggerMode.is_valid(self._mode):
      raise ValueError("Logger mode %s is invalid." % self._mode)

  def get_handlers(self):
    """
    Creates stdout/stderr handler by provided configuration
    """
    return (self._get_handler(self.STDOUT),
            self._get_handler(self.STDERR),
            self._handlers_are_files())

  def _handlers_are_files(self):
    """
    Returns True if both the handlers are standard file objects.
    """
    return (self._destination == LoggerDestination.CONSOLE or
        (self._destination == LoggerDestination.FILE and self._mode == LoggerMode.STANDARD))

  def _get_handler(self, name):
    """
    Constructs correct handler or file object based on the provided configuration.
    """

    # On no destination write logs to /dev/null
    if self._destination == LoggerDestination.NONE:
      return StreamHandler(safe_open(os.devnull, 'w'))

    # Streamed logs to predefined outputs
    if self._destination == LoggerDestination.CONSOLE:
      return sys.stdout if name == self.STDOUT else sys.stderr

    # Streaming AND file logs are required
    if self._destination == LoggerDestination.BOTH:
      return TeeHandler(self._get_stream(name), self._get_file(name))

    # File only logs are required
    return self._get_file(name)

  def _get_file(self, name):
    if self._mode == LoggerMode.STANDARD:
      return safe_open(self._get_log_path(name), mode='a')
    if self._mode == LoggerMode.ROTATE:
      log_size = int(self._rotate_log_size.as_(Data.BYTES))
      return RotatingFileHandler(self._get_log_path(name),
                                 log_size,
                                 self._rotate_log_backups)

  def _get_stream(self, name):
    """
    Returns OS stream by name
    """
    if name == self.STDOUT:
      return StreamHandler(sys.stdout)
    if name == self.STDERR:
      return StreamHandler(sys.stderr)

  def _get_log_path(self, log_name):
    return self._pathspec.with_filename(log_name).getpath('process_logdir')


class RotatingFileHandler(object):
  """
  File handler that implements max size/rotation.
  """

  def __init__(self, filename, max_bytes, max_backups, mode='w'):
    """
      required:
        filename    = The file name.
        max_bytes   = The maximum size of an individual log file.
        max_backups = The maximum number of log file backups to create.

      optional:
        mode = Mode to open the file in.
    """
    if max_bytes > 0 and max_backups <= 0:
      raise ValueError('A positive value for max_backups must be specified if max_bytes > 0.')
    self._max_bytes = max_bytes
    self._max_backups = max_backups
    self.file = safe_open(filename, mode=mode)
    self.filename = filename
    self.mode = mode
    self.closed = False

  def close(self):
    if not self.closed:
      self.file.close()
      self.closed = True

  def write(self, b):
    self.file.write(b)
    self.file.flush()
    if self.should_rollover():
      self.rollover()

  def swap_files(self, src, tgt):
    if os.path.exists(tgt):
      safe_delete(tgt)

    try:
      os.rename(src, tgt)
    except OSError as e:
      if e.errno != errno.ENOENT:
        raise

  def make_indexed_filename(self, index):
    return '%s.%d' % (self.filename, index)

  def should_rollover(self):
    if self._max_bytes <= 0 or self._max_backups <= 0:
      return False

    if self.file.tell() >= self._max_bytes:
      return True

    return False

  def rollover(self):
    """
    Perform the rollover of the log.
    """
    self.file.close()
    for i in range(self._max_backups - 1, 0, -1):
      src = self.make_indexed_filename(i)
      tgt = self.make_indexed_filename(i + 1)
      if os.path.exists(src):
        self.swap_files(src, tgt)

    self.swap_files(self.filename, self.make_indexed_filename(1))
    self.file = safe_open(self.filename, mode='w')


class StreamHandler(object):
  """
  Stream handler wraps stream objects and allows configuration of whether objects
  should be closed when ending a subprocess.
  """

  def __init__(self, stream, close=False):
    """
    stream = Wrapped stream object.
    """
    self._stream = stream
    self._close = close

  def write(self, b):
    self._stream.write(b)
    self._stream.flush()

  def close(self):
    if self._close:
      self._stream.close()


class TeeHandler(object):
  """
  Tee handler mimicks the unix tee command and splits output between two destinations
  """

  def __init__(self, first, second):
    """
      required:
        first  = First destination
        second = Second destination
    """
    self._first = first
    self._second = second

  def write(self, b):
    self._first.write(b)
    self._second.write(b)

  def close(self):
    self._first.close()
    self._second.close()
