示例#1
0
 def _http_get(self, path: str) -> SpooledTemporaryFile:
     url = urllib.parse.urljoin(self.repo_url.rstrip("/") + "/", path)
     logger.debug("Fetching %s", url)
     response = self._session.get(url, headers=HEADERS)
     buffer = SpooledTemporaryFile(max_size=100 * 1024 * 1024)
     for chunk in response.iter_content(chunk_size=10 * 1024 * 1024):
         buffer.write(chunk)
     buffer.flush()
     buffer.seek(0)
     return buffer
示例#2
0
class RemoteFileBuffer(object):
    """File-like object providing buffer for local file operations.

    Instances of this class manage a local tempfile buffer corresponding
    to the contents of a remote file.  All reads and writes happen locally,
    with the content being copied to the remote file only on flush() or
    close().

    Instances of this class are returned by S3FS.open, but it is desgined
    to be usable by any FS subclass that manages remote files.
    """

    def __init__(self,fs,path,mode):
        self.file = TempFile()
        self.fs = fs
        self.path = path
        self.mode = mode

    def __del__(self):
        if not self.closed:
            self.close()

    #  This is lifted straight from the stdlib's tempfile.py
    def __getattr__(self,name):
        file = self.__dict__['file']
        a = getattr(file, name)
        if not issubclass(type(a), type(0)):
            setattr(self, name, a)
        return a

    def __enter__(self):
        self.file.__enter__()
        return self

    def __exit__(self,exc,value,tb):
        self.close()
        return False

    def __iter__(self):
        return iter(self.file)

    def flush(self):
        self.file.flush()
        if "w" in self.mode or "a" in self.mode or "+" in self.mode:
            pos = self.file.tell()
            self.file.seek(0)
            self.fs.setcontents(self.path,self.file)
            self.file.seek(pos)

    def close(self):
        if "w" in self.mode or "a" in self.mode or "+" in self.mode:
            self.file.seek(0)
            self.fs.setcontents(self.path,self.file)
        self.file.close()
示例#3
0
    def fetch_pack_from_origin(
        self,
        origin_url: str,
        base_repo: RepoRepresentation,
        do_activity: Callable[[bytes], None],
    ) -> FetchPackReturn:
        """Fetch a pack from the origin"""

        pack_buffer = SpooledTemporaryFile(max_size=self.temp_file_cutoff)
        transport_url = origin_url

        logger.debug("Transport url to communicate with server: %s", transport_url)

        client, path = dulwich.client.get_transport_and_path(
            transport_url, thin_packs=False
        )

        logger.debug("Client %s to fetch pack at %s", client, path)

        size_limit = self.pack_size_bytes

        def do_pack(data: bytes) -> None:
            cur_size = pack_buffer.tell()
            would_write = len(data)
            if cur_size + would_write > size_limit:
                raise IOError(
                    f"Pack file too big for repository {origin_url}, "
                    f"limit is {size_limit} bytes, current size is {cur_size}, "
                    f"would write {would_write}"
                )

            pack_buffer.write(data)

        pack_result = client.fetch_pack(
            path,
            base_repo.determine_wants,
            base_repo.graph_walker(),
            do_pack,
            progress=do_activity,
        )

        remote_refs = pack_result.refs or {}
        symbolic_refs = pack_result.symrefs or {}

        pack_buffer.flush()
        pack_size = pack_buffer.tell()
        pack_buffer.seek(0)

        logger.debug("fetched_pack_size=%s", pack_size)

        # check if repository only supports git dumb transfer protocol,
        # fetched pack file will be empty in that case as dulwich do
        # not support it and do not fetch any refs
        self.dumb = transport_url.startswith("http") and getattr(client, "dumb", False)

        return FetchPackReturn(
            remote_refs=utils.filter_refs(remote_refs),
            symbolic_refs=utils.filter_refs(symbolic_refs),
            pack_buffer=pack_buffer,
            pack_size=pack_size,
        )
    def fetch_media(self, url, partial_fetch=False):
        """Retrieves a given media object from a remote (HTTP) location
        and returns the content-type and a file-like object containing
        the media content.

        The file-like object is a temporary file that - depending on the
        size - lives in memory or on disk. Once the file is closed, the
        contents are removed from storage.

        :param url: the URL of the media asset.
        :type url: str.
        :param partial_fetch: determines if the the complete file should
            be fetched, or if only the first 2 MB should be retrieved.
            This feature is used to prevent complete retrieval of large
            a/v material.
        :type partial_fetch: bool.
        :returns: a tuple with the ``content-type``, ``content-lenght``
            and a file-like object containing the media content. The
            value of ``content-length`` will be ``None`` in case
            a partial fetch is requested and ``content-length`` is not
            returned by the remote server.
        """

        http_resp = self.http_session.get(url, stream=True, timeout=(60, 120))
        http_resp.raise_for_status()

        if not os.path.exists(TEMP_DIR_PATH):
            log.debug('Creating temp directory %s' % TEMP_DIR_PATH)
            os.makedirs(TEMP_DIR_PATH)

        # Create a temporary file to store the media item, write the file
        # to disk if it is larger than 1 MB.
        media_file = SpooledTemporaryFile(max_size=1024 * 1024,
                                          prefix='oad_m_',
                                          suffix='.tmp',
                                          dir=TEMP_DIR_PATH)

        # When a partial fetch is requested, request up to two MB
        partial_target_size = 1024 * 1024 * 2
        content_length = http_resp.headers.get('content-length')
        if content_length and int(content_length) < partial_target_size:
            partial_target_size = int(content_length)

        retrieved_bytes = 0
        for chunk in http_resp.iter_content(chunk_size=512 * 1024):
            if chunk:  # filter out keep-alive chunks
                media_file.write(chunk)
                retrieved_bytes += len(chunk)

            if partial_fetch and retrieved_bytes >= partial_target_size:
                break

        media_file.flush()
        log.debug('Fetched media item %s [%s/%s]' %
                  (url, retrieved_bytes, content_length))

        # If the server doens't provide a content-length and this isn't
        # a partial fetch, determine the size by looking at the retrieved
        # content
        if not content_length and not partial_fetch:
            media_file.seek(0, 2)
            content_length = media_file.tell()

        return (http_resp.headers.get('content-type'), content_length,
                media_file)
示例#5
0
class RecvWindow(object):
    """
    Receiver windows receive data on a socket to fill their
    buffer.

    A window instance has the following attributes:

        name        Name of the window. The name of all windows is shown by
                    '> win show' in command line interface. It's usually the
                    name of the plugin appended by the arguments passed for
                    execution (e.g. ps.plg -A -t curesec.com)
        status      indicates the state of the plugin (e.g. RUNNING or FINISHED)
        req_t       request tuple that triggered plugin execution. It's used to
                    verify incoming packets
        plg_id      id of the plugin
        rid         request id of the packet triggering execution
        wid         the id of the window. Derived from a class variable which
                    is increased after window creation. The window id is not
                    fixed. So every new client session, the windows might have
                    a new id.
        session     the actual client session
        is_used     flag, that indicates, whether the windows is used by user
                    (via '> win show <id>')

    """

    #: command to close window
    CMD_CLOSE = "close"

    #: window id, increased for every new window
    ID = 1

    def __init__(self, name, req_t, session):
        logger.debug("creating window %s", name)
        self.name = name
        self.status = RUNNING
        self.req_t = req_t
        self.sid = req_t[1]
        self.rid = req_t[2]
        self.plg_id = req_t[4]
        self.wid = RecvWindow.ID
        self.session = session
        self.is_used = False
        self.closed = False

        self._buffer = SpooledTemporaryFile(max_size="8096", mode="wr+b")
        RecvWindow.ID += 1

    def close(self):
        self.closed = True
        print("closing window %d" % self.wid)

    def set_state(self, input_code):
        state = None
        try:
            state = STATE[input_code + 1]
        except (IndexError, TypeError) as e:
            logger.warning('Could not get window state for code %d: %s',
                           input_code, e)
            logger.exception(e)

        if not state:
            state = "Unknown"

        if self.is_used and state in (FINISHED, ERROR, KILLED):
            self.is_used = False

        code = "(%s)" % str(input_code)
        self.status = "{state} {code}".format(state=state, code=code)

        logger.debug("new window(id:%s) state: %s", self.wid, self.status)

    def _write_to_output(self, data):
        # if no more data is coming, set window state to
        # finished and break if we have a non-interactive
        # window (otherwise we need to wait for the user
        # to enter further input).
        #if not data:
        #    self.status = FINISHED
        # TODO: that seems a bit heuristic to me ...

        # ok, we have data, so verify 'em and print it. this is
        # done by unpack
        try:
            actual_output = self._unpack(data)
            if actual_output:
                sys.stdout.write(actual_output)
                sys.stdout.flush()
                self._buffer.write(actual_output)
                self._buffer.flush()

        except errors.ClientError as e:
            if e.request:
                logger.warning('Request: %s', conn.ccdlib.pkt2str(e.request))
            if e.response:
                logger.warning('Response: %s', conn.ccdlib.pkt2str(e.response))
            logger.debug("received invalid packet:'%s'", e)
            logger.exception(e)
        except Exception as e:
            logger.debug("received invalid packet:'%s'", e)
            logger.exception(e)

    def _unpack(self, resp_t):
        #if not resp_t[0] == cnnm.ccdlib.OP_SUCCESS:
        #    raise Exception("received packet that indicates a failure(%s)!",
        #                    hex(resp_t[0]))

        ## sid
        #if not resp_t[1] == self.req_t[1]:
        #    raise Exception("received packet with invalid session id(%s)!",
        #                    resp_t[1])

        ## rid
        #if not resp_t[2] == self.req_t[2]:
        #    raise Exception("received packet with invalid rid(%s)!",
        #                    resp_t[2])

        ## plg_id
        #if not resp_t[4] == self.req_t[4]:
        #    raise Exception("received packet with invalid plugin id(%s)!",
        #                    resp_t[4])

        # method indicates a plugin's state change.
        if resp_t[6] == conn.ccdlib.MTH_TERMINATE:
            state = 2
            try:
                state = resp_t[-1]["code"]
            except (KeyError, IndexError):
                raise errors.InvalidServerResponse(
                    message="No state in terminate response payload",
                    request=self.req_t,
                    response=resp_t)
            self.set_state(state)
            return

        elif resp_t[6] == conn.ccdlib.MTH_INPUT:
            return ""

        # plugin gives outpout
        elif not resp_t[6] == conn.ccdlib.MTH_OUTPUT:
            try:
                MTH_str = conn.ccdlib.rev_MTH_map[resp_t[6]]
            except:
                MTH_str = repr(resp_t[6])
            raise errors.InvalidServerResponse(
                message=
                "Invalid method %s, expected MTH_OUTPUT, _TERMINATE or _INPUT"
                % MTH_str,
                request=self.req_t,
                response=resp_t)

        return resp_t[-1]

    def _pack(self, data):
        req_t = list(self.req_t)

        req_t[
            0] = conn.ccdlib.OP_PLUGIN  # indicate operation concerning plugin
        req_t[6] = conn.ccdlib.MTH_INPUT  # indicate plugin input is coming
        req_t[-1] = data  # payload

        # if no input to commit, raise EmptyException
        if not req_t[-1]:
            raise conn.ccdlib.EmptyException

        return tuple(req_t)

    def use(self):
        # to get the plugin's output, we need to register
        rid = self.req_t[2]
        plg_id = self.req_t[4]
        pld = comm.sendpacket(self.session,
                              op=conn.ccdlib.OP_REGISTER,
                              plg=plg_id,
                              pld=dict(rid=rid))[-1]
        interactive = pld["interactive"]
        logger.debug("successfully registered to plugin. have fun!")

        if interactive:
            print("Enter '%s' to leave interactive mode." %
                  RecvWindow.CMD_CLOSE)

        sock = self.session.sock

        # write content that is already buffered
        for line in self._buffer:
            sys.stdout.write(line)

        # start window interaction
        self.is_used = True
        while self.is_used:

            try:
                r, _, _ = select.select([sys.stdin, sock], [], [])
            except KeyboardInterrupt:
                print("catched keyboard interrupt. closing window.")
                self.is_used = False
                break

            # if there is something to print to window
            if sock in r:
                # read data, that should be written to stdout.
                # the data is sent by the ccd
                resp_t = None
                try:
                    resp_t = conn.ccdlib.recv(sock)
                    #resp_t = comm.wait_for_response(self.session, self.req_t)
                except Exception as e:
                    logger.error("Exception while receiving data: %s", e)
                    logger.exception(e)
                    self.is_used = False
                    break

                if resp_t:
                    self._write_to_output(resp_t)

            # there is something to read form window console
            if sys.stdin in r:
                if not interactive:
                    # any keystroke sends active plugin to background
                    self.is_used = False
                    break

                # read the user's input to be sent to the ccd plugin
                #data = sys.stdin.readline()
                data = ''
                try:
                    data = os.read(sys.stdin.fileno(), 1024)
                except Exception, e:
                    logger.error("Error reading from stdin: '%s'!", e)
                    logger.exception(e)
                    continue

                logger.debug('all read on stdin \'%s\'', data)

                # catch command to close the interactive window
                if data.rstrip() == RecvWindow.CMD_CLOSE:
                    self.is_used = False
                    break

                # the actual write of the user's input
                actual_packet = self._pack(data)

                try:
                    conn.ccdlib.send(sock, actual_packet)
                    #comm.wait_for_response(self.session, actual_packet)
                except Exception as e:
                    logger.error("Failed to send user input to plugin!:'%s'",
                                 e)
                    logger.exception(e)

        try:
            comm.sendpacket(self.session,
                            op=conn.ccdlib.OP_UNREGISTER,
                            plg=plg_id,
                            pld=dict(rid=rid))
        except KeyboardInterrupt:
            print("catched keyboard interrupt. closing window.")
        except Exception as e:
            logger.error("Failed to unregister properly:'%s'!", e)
            logger.exception(e)
    def flush(self):
        """Flush the data to the server."""

        SpooledTemporaryFile.flush(self)
        self.commit()
示例#7
0
文件: client.py 项目: Alexis-D/DFS
    def flush(self):
        """Flush the data to the server."""

        SpooledTemporaryFile.flush(self)
        self.commit()
示例#8
0
 def flush(self): # flush data to server
 
     SpooledTemporaryFile.flush(self)
     self.commit()
示例#9
0
    def fetch_media(self, url, partial_fetch=False):
        """Retrieves a given media object from a remote (HTTP) location
        and returns the content-type and a file-like object containing
        the media content.

        The file-like object is a temporary file that - depending on the
        size - lives in memory or on disk. Once the file is closed, the
        contents are removed from storage.

        :param url: the URL of the media asset.
        :type url: str.
        :param partial_fetch: determines if the the complete file should
            be fetched, or if only the first 2 MB should be retrieved.
            This feature is used to prevent complete retrieval of large
            a/v material.
        :type partial_fetch: bool.
        :returns: a tuple with the ``content-type``, ``content-lenght``
            and a file-like object containing the media content. The
            value of ``content-length`` will be ``None`` in case
            a partial fetch is requested and ``content-length`` is not
            returned by the remote server.
        """

        http_resp = self.http_session.get(url, stream=True, timeout=(60, 120))
        http_resp.raise_for_status()

        if not os.path.exists(TEMP_DIR_PATH):
            log.debug('Creating temp directory %s' % TEMP_DIR_PATH)
            os.makedirs(TEMP_DIR_PATH)

        # Create a temporary file to store the media item, write the file
        # to disk if it is larger than 1 MB.
        media_file = SpooledTemporaryFile(max_size=1024*1024, prefix='ocd_m_',
                                          suffix='.tmp',
                                          dir=TEMP_DIR_PATH)

        # When a partial fetch is requested, request up to two MB
        partial_target_size = 1024*1024*2
        content_length = http_resp.headers.get('content-length')
        if content_length and int(content_length) < partial_target_size:
            partial_target_size = int(content_length)

        retrieved_bytes = 0
        for chunk in http_resp.iter_content(chunk_size=512*1024):
            if chunk:  # filter out keep-alive chunks
                media_file.write(chunk)
                retrieved_bytes += len(chunk)

            if partial_fetch and retrieved_bytes >= partial_target_size:
                break

        media_file.flush()
        log.debug('Fetched media item %s [%s/%s]' % (url, retrieved_bytes,
                                                     content_length))

        # If the server doens't provide a content-length and this isn't
        # a partial fetch, determine the size by looking at the retrieved
        # content
        if not content_length and not partial_fetch:
            media_file.seek(0, 2)
            content_length = media_file.tell()

        return (
            http_resp.headers.get('content-type'),
            content_length,
            media_file
        )
示例#10
0
class Buffer(FileWrapper):
    """Class implementing buffereing of input and output streams.
    
    This class uses a separate buffer file to hold the contents of the
    underlying file while they are being manipulated.  As data is read
    it is duplicated into the buffer, and data is written from the buffer
    back to the file on close.
    """
    
    def __init__(self,fileobj,mode=None,max_size_in_memory=1024*8):
        """Buffered file wrapper constructor."""
        self._buffer = SpooledTemporaryFile(max_size=max_size_in_memory)
        self._in_eof = False
        self._in_pos = 0
        super(Buffer,self).__init__(fileobj,mode)

    def _buffer_chunks(self):
        chunk = self._buffer.read(16*1024)
        if chunk == "":
            yield chunk
        else:
            while chunk != "":
                yield chunk
                chunk = self._buffer.read(16*1024)

    def _write_out_buffer(self):
        if self._check_mode("r"):
            self._read_rest()
            if "a" in self.mode:
                self._buffer.seek(self._in_pos)
                self._fileobj.seek(self._in_pos)
            else:
                self._fileobj.seek(0)
                self._buffer.seek(0)
        else:
            self._buffer.seek(0)
        for chunk in self._buffer_chunks():
            self._fileobj.write(chunk)
 
    def flush(self):
        # flush the buffer; we only write to the underlying file on close
        self._buffer.flush()

    def close(self):
        if self.closed:
            return
        if self._check_mode("w"):
            self._write_out_buffer()
        super(Buffer,self).close()
        self._buffer.close()

    def _read(self,sizehint=-1):
        #  First return any data available from the buffer.
        #  Since we don't flush the buffer after every write, certain OSes
        #  (guess which!) will happy read junk data from the end of it.
        #  Instead, we explicitly read only up to self._in_pos.
        if not self._in_eof:
            buffered_size = self._in_pos - self._buffer.tell()
            if sizehint >= 0:
                buffered_size = min(sizehint,buffered_size)
        else:
            buffered_size = sizehint
        data = self._buffer.read(buffered_size)
        if data != "":
            return data
        # Then look for more data in the underlying file
        if self._in_eof:
            return None
        data = self._fileobj.read(sizehint)
        self._in_pos += len(data)
        self._buffer.write(data)
        if sizehint < 0 or len(data) < sizehint:
            self._in_eof = True
            self._buffer.flush()
        return data

    def _write(self,data,flushing=False):
        self._buffer.write(data)
        if self._check_mode("r") and not self._in_eof:
            diff = self._buffer.tell() - self._in_pos
            if diff > 0:
                junk = self._fileobj.read(diff)
                self._in_pos += len(junk)
                if len(junk) < diff:
                    self._in_eof = True
                    self._buffer.flush()
    
    def _seek(self,offset,whence):
        # Ensure we've read enough to simply do the seek on the buffer
        if self._check_mode("r") and not self._in_eof:
            if whence == 0:
                if offset > self._in_pos:
                    self._read_rest()
            if whence == 1:
                if self._buffer.tell() + offset > self._in_pos:
                    self._read_rest()
            if whence == 2:
                self._read_rest()
        # Then just do it on the buffer...
        self._buffer.seek(offset,whence)

    def _tell(self):
        return self._buffer.tell()
        
    def _read_rest(self):
        """Read the rest of the input stream."""
        if self._in_eof:
            return
        pos = self._buffer.tell()
        self._buffer.seek(0,2)
        data = self._fileobj.read(self._bufsize)
        while data:
            self._in_pos += len(data)
            self._buffer.write(data)
            data = self._fileobj.read(self._bufsize)
        self._in_eof = True 
        self._buffer.flush()
        self._buffer.seek(pos)
示例#11
0
class Buffer(FileWrapper):
    """Class implementing buffering of input and output streams.
    
    This class uses a separate buffer file to hold the contents of the
    underlying file while they are being manipulated.  As data is read
    it is duplicated into the buffer, and data is written from the buffer
    back to the file on close.
    """
    def __init__(self, fileobj, mode=None, max_size_in_memory=1024 * 8):
        """Buffered file wrapper constructor."""
        self._buffer = SpooledTemporaryFile(max_size=max_size_in_memory)
        self._in_eof = False
        self._in_pos = 0
        self._was_truncated = False
        super(Buffer, self).__init__(fileobj, mode)

    def _buffer_size(self):
        try:
            return len(self._buffer.file.getvalue())
        except AttributeError:
            return os.fstat(self._buffer.fileno()).st_size

    def _buffer_chunks(self):
        chunk = self._buffer.read(16 * 1024)
        if chunk == "":
            yield chunk
        else:
            while chunk != "":
                yield chunk
                chunk = self._buffer.read(16 * 1024)

    def _write_out_buffer(self):
        if self._check_mode("r"):
            self._read_rest()
            if "a" in self.mode:
                self._buffer.seek(self._in_pos)
                self._fileobj.seek(self._in_pos)
            else:
                self._fileobj.seek(0)
                self._buffer.seek(0)
        else:
            self._buffer.seek(0)
        if self._was_truncated:
            self._fileobj.truncate(0)
            self._was_truncated = False
        for chunk in self._buffer_chunks():
            self._fileobj.write(chunk)

    def flush(self):
        # flush the buffer; we only write to the underlying file on close
        self._buffer.flush()

    def close(self):
        if self.closed:
            return
        if self._check_mode("w"):
            self._write_out_buffer()
        super(Buffer, self).close()
        self._buffer.close()

    def _read(self, sizehint=-1):
        #  First return any data available from the buffer.
        #  Since we don't flush the buffer after every write, certain OSes
        #  (guess which!) will happily read junk data from the end of it.
        #  Instead, we explicitly read only up to self._in_pos.
        if not self._in_eof:
            buffered_size = self._in_pos - self._buffer.tell()
            if sizehint >= 0:
                buffered_size = min(sizehint, buffered_size)
        else:
            buffered_size = sizehint
        data = self._buffer.read(buffered_size)
        if data != "":
            return data
        # Then look for more data in the underlying file
        if self._in_eof:
            return None
        data = self._fileobj.read(sizehint)
        self._in_pos += len(data)
        self._buffer.write(data)
        if sizehint < 0 or len(data) < sizehint:
            self._in_eof = True
            self._buffer.flush()
        return data

    def _write(self, data, flushing=False):
        self._buffer.write(data)
        if self._check_mode("r") and not self._in_eof:
            diff = self._buffer.tell() - self._in_pos
            if diff > 0:
                junk = self._fileobj.read(diff)
                self._in_pos += len(junk)
                if len(junk) < diff:
                    self._in_eof = True
                    self._buffer.flush()

    def _seek(self, offset, whence):
        # Ensure we've read enough to simply do the seek on the buffer
        if self._check_mode("r") and not self._in_eof:
            if whence == 0:
                if offset > self._in_pos:
                    self._read_rest()
            if whence == 1:
                if self._buffer.tell() + offset > self._in_pos:
                    self._read_rest()
            if whence == 2:
                self._read_rest()
        # Then just do it on the buffer...
        self._buffer.seek(offset, whence)

    def _tell(self):
        return self._buffer.tell()

    def _truncate(self, size):
        if self._check_mode("r") and not self._in_eof:
            if size > self._in_pos:
                self._read_rest()
        self._in_eof = True
        try:
            self._buffer.truncate(size)
        except TypeError:
            et, ev, tb = sys.exc_info()
            # SpooledTemporaryFile.truncate() doesn't accept size paramter.
            try:
                self._buffer._file.truncate(size)
            except Exception:
                raise et, ev, tb
        # StringIO objects don't truncate to larger size correctly.
        if hasattr(self._buffer, "_file"):
            _file = self._buffer._file
            if hasattr(_file, "getvalue"):
                if len(_file.getvalue()) != size:
                    curpos = _file.tell()
                    _file.seek(0, 2)
                    _file.write("\x00" * (size - len(_file.getvalue())))
                    _file.seek(curpos)
        self._was_truncated = True

    def _read_rest(self):
        """Read the rest of the input stream."""
        if self._in_eof:
            return
        pos = self._buffer.tell()
        self._buffer.seek(0, 2)
        data = self._fileobj.read(self._bufsize)
        while data:
            self._in_pos += len(data)
            self._buffer.write(data)
            data = self._fileobj.read(self._bufsize)
        self._in_eof = True
        self._buffer.flush()
        self._buffer.seek(pos)
    def flush(self):

        SpooledTemporaryFile.flush(self)
        self.commit()
class RecvWindow(object):
    """
    Receiver windows receive data on a socket to fill their
    buffer.

    A window instance has the following attributes:

        name        Name of the window. The name of all windows is shown by
                    '> win show' in command line interface. It's usually the
                    name of the plugin appended by the arguments passed for
                    execution (e.g. ps.plg -A -t curesec.com)
        status      indicates the state of the plugin (e.g. RUNNING or FINISHED)
        req_t       request tuple that triggered plugin execution. It's used to
                    verify incoming packets
        plg_id      id of the plugin
        rid         request id of the packet triggering execution
        wid         the id of the window. Derived from a class variable which
                    is increased after window creation. The window id is not
                    fixed. So every new client session, the windows might have
                    a new id.
        session     the actual client session
        is_used     flag, that indicates, whether the windows is used by user
                    (via '> win show <id>')

    """

    #: command to close window
    CMD_CLOSE = "close"

    #: window id, increased for every new window
    ID = 1

    def __init__(self, name, req_t, session):
        logger.debug("creating window %s", name)
        self.name = name
        self.status = RUNNING
        self.req_t = req_t
        self.sid = req_t[1]
        self.rid = req_t[2]
        self.plg_id = req_t[4]
        self.wid = RecvWindow.ID
        self.session = session
        self.is_used = False
        self.closed = False

        self._buffer = SpooledTemporaryFile(max_size="8096", mode="wr+b")
        RecvWindow.ID += 1

    def close(self):
        self.closed = True
        print("closing window %d" % self.wid)

    def set_state(self, input_code):
        state = None
        try:
            state = STATE[input_code + 1]
        except (IndexError, TypeError) as e:
            logger.warning('Could not get window state for code %d: %s', input_code, e)
            logger.exception(e)

        if not state:
            state = "Unknown"

        if self.is_used and state in (FINISHED, ERROR, KILLED):
            self.is_used = False

        code = "(%s)" % str(input_code)
        self.status = "{state} {code}".format(state=state, code=code)

        logger.debug("new window(id:%s) state: %s", self.wid, self.status)

    def _write_to_output(self, data):
        # if no more data is coming, set window state to
        # finished and break if we have a non-interactive
        # window (otherwise we need to wait for the user
        # to enter further input).
        #if not data:
        #    self.status = FINISHED
        # TODO: that seems a bit heuristic to me ...

        # ok, we have data, so verify 'em and print it. this is
        # done by unpack
        try:
            actual_output = self._unpack(data)
            if actual_output:
                sys.stdout.write(actual_output)
                sys.stdout.flush()
                self._buffer.write(actual_output)
                self._buffer.flush()

        except errors.ClientError as e:
            if e.request:
                logger.warning('Request: %s', conn.ccdlib.pkt2str(e.request))
            if e.response:
                logger.warning('Response: %s', conn.ccdlib.pkt2str(e.response))
            logger.debug("received invalid packet:'%s'", e)
            logger.exception(e)
        except Exception as e:
            logger.debug("received invalid packet:'%s'", e)
            logger.exception(e)

    def _unpack(self, resp_t):
        #if not resp_t[0] == cnnm.ccdlib.OP_SUCCESS:
        #    raise Exception("received packet that indicates a failure(%s)!",
        #                    hex(resp_t[0]))

        ## sid
        #if not resp_t[1] == self.req_t[1]:
        #    raise Exception("received packet with invalid session id(%s)!",
        #                    resp_t[1])

        ## rid
        #if not resp_t[2] == self.req_t[2]:
        #    raise Exception("received packet with invalid rid(%s)!",
        #                    resp_t[2])

        ## plg_id
        #if not resp_t[4] == self.req_t[4]:
        #    raise Exception("received packet with invalid plugin id(%s)!",
        #                    resp_t[4])

        # method indicates a plugin's state change.
        if resp_t[6] == conn.ccdlib.MTH_TERMINATE:
            state = 2
            try:
                state = resp_t[-1]["code"]
            except (KeyError, IndexError):
                raise errors.InvalidServerResponse(
                    message="No state in terminate response payload",
                    request=self.req_t, response=resp_t)
            self.set_state(state)
            return

        elif resp_t[6] == conn.ccdlib.MTH_INPUT:
            return ""

        # plugin gives outpout
        elif not resp_t[6] == conn.ccdlib.MTH_OUTPUT:
            try:
                MTH_str = conn.ccdlib.rev_MTH_map[resp_t[6]]
            except:
                MTH_str = repr(resp_t[6])
            raise errors.InvalidServerResponse(
                message="Invalid method %s, expected MTH_OUTPUT, _TERMINATE or _INPUT" % MTH_str,
                request=self.req_t, response=resp_t)

        return resp_t[-1]

    def _pack(self, data):
        req_t = list(self.req_t)

        req_t[0] = conn.ccdlib.OP_PLUGIN  # indicate operation concerning plugin
        req_t[6] = conn.ccdlib.MTH_INPUT  # indicate plugin input is coming
        req_t[-1] = data  # payload

        # if no input to commit, raise EmptyException
        if not req_t[-1]:
            raise conn.ccdlib.EmptyException

        return tuple(req_t)

    def use(self):
        # to get the plugin's output, we need to register
        rid = self.req_t[2]
        plg_id = self.req_t[4]
        pld = comm.sendpacket(self.session,
                              op=conn.ccdlib.OP_REGISTER,
                              plg=plg_id,
                              pld=dict(rid=rid))[-1]
        interactive = pld["interactive"]
        logger.debug("successfully registered to plugin. have fun!")

        if interactive:
            print("Enter '%s' to leave interactive mode." %
                  RecvWindow.CMD_CLOSE)

        sock = self.session.sock

        # write content that is already buffered
        for line in self._buffer:
            sys.stdout.write(line)

        # start window interaction
        self.is_used = True
        while self.is_used:

            try:
                r, _, _ = select.select([sys.stdin, sock], [], [])
            except KeyboardInterrupt:
                print("catched keyboard interrupt. closing window.")
                self.is_used = False
                break

            # if there is something to print to window
            if sock in r:
                # read data, that should be written to stdout.
                # the data is sent by the ccd
                resp_t = None
                try:
                    resp_t = conn.ccdlib.recv(sock)
                    #resp_t = comm.wait_for_response(self.session, self.req_t)
                except Exception as e:
                    logger.error("Exception while receiving data: %s", e)
                    logger.exception(e)
                    self.is_used = False
                    break

                if resp_t:
                    self._write_to_output(resp_t)

            # there is something to read form window console
            if sys.stdin in r:
                if not interactive:
                    # any keystroke sends active plugin to background
                    self.is_used = False
                    break

                # read the user's input to be sent to the ccd plugin
                #data = sys.stdin.readline()
                data = ''
                try:
                    data = os.read(sys.stdin.fileno(), 1024)
                except Exception, e:
                    logger.error("Error reading from stdin: '%s'!", e)
                    logger.exception(e)
                    continue

                logger.debug('all read on stdin \'%s\'', data)

                # catch command to close the interactive window
                if data.rstrip() == RecvWindow.CMD_CLOSE:
                    self.is_used = False
                    break

                # the actual write of the user's input
                actual_packet = self._pack(data)

                try:
                    conn.ccdlib.send(sock, actual_packet)
                    #comm.wait_for_response(self.session, actual_packet)
                except Exception as e:
                    logger.error("Failed to send user input to plugin!:'%s'", e)
                    logger.exception(e)

        try:
            comm.sendpacket(self.session,
                            op=conn.ccdlib.OP_UNREGISTER,
                            plg=plg_id,
                            pld=dict(rid=rid))
        except KeyboardInterrupt:
                print("catched keyboard interrupt. closing window.")
        except Exception as e:
            logger.error("Failed to unregister properly:'%s'!", e)
            logger.exception(e)
示例#14
0
 def exposed_cache_report(self, exporter_id, options):
     exporter = self._sa._manager.create_exporter(exporter_id, options)
     exported_stream = SpooledTemporaryFile(max_size=102400, mode='w+b')
     exporter.dump_from_session_manager(self._sa._manager, exported_stream)
     exported_stream.flush()
     return exported_stream
示例#15
0
class Buffer(FileWrapper):
    """Class implementing buffering of input and output streams.
    
    This class uses a separate buffer file to hold the contents of the
    underlying file while they are being manipulated.  As data is read
    it is duplicated into the buffer, and data is written from the buffer
    back to the file on close.
    """

    def __init__(self, fileobj, mode=None, max_size_in_memory=1024 * 8):
        """Buffered file wrapper constructor."""
        self._buffer = SpooledTemporaryFile(max_size=max_size_in_memory)
        self._in_eof = False
        self._in_pos = 0
        self._was_truncated = False
        super(Buffer, self).__init__(fileobj, mode)

    def _buffer_size(self):
        try:
            return len(self._buffer.file.getvalue())
        except AttributeError:
            return os.fstat(self._buffer.fileno()).st_size

    def _buffer_chunks(self):
        chunk = self._buffer.read(16 * 1024)
        if chunk == "":
            yield chunk
        else:
            while chunk != "":
                yield chunk
                chunk = self._buffer.read(16 * 1024)

    def _write_out_buffer(self):
        if self._check_mode("r"):
            self._read_rest()
            if "a" in self.mode:
                self._buffer.seek(self._in_pos)
                self._fileobj.seek(self._in_pos)
            else:
                self._fileobj.seek(0)
                self._buffer.seek(0)
        else:
            self._buffer.seek(0)
        if self._was_truncated:
            self._fileobj.truncate(0)
            self._was_truncated = False
        for chunk in self._buffer_chunks():
            self._fileobj.write(chunk)

    def flush(self):
        # flush the buffer; we only write to the underlying file on close
        self._buffer.flush()

    def close(self):
        if self.closed:
            return
        if self._check_mode("w"):
            self._write_out_buffer()
        super(Buffer, self).close()
        self._buffer.close()

    def _read(self, sizehint=-1):
        #  First return any data available from the buffer.
        #  Since we don't flush the buffer after every write, certain OSes
        #  (guess which!) will happily read junk data from the end of it.
        #  Instead, we explicitly read only up to self._in_pos.
        if not self._in_eof:
            buffered_size = self._in_pos - self._buffer.tell()
            if sizehint >= 0:
                buffered_size = min(sizehint, buffered_size)
        else:
            buffered_size = sizehint
        data = self._buffer.read(buffered_size)
        if data != "":
            return data
        # Then look for more data in the underlying file
        if self._in_eof:
            return None
        data = self._fileobj.read(sizehint)
        self._in_pos += len(data)
        self._buffer.write(data)
        if sizehint < 0 or len(data) < sizehint:
            self._in_eof = True
            self._buffer.flush()
        return data

    def _write(self, data, flushing=False):
        self._buffer.write(data)
        if self._check_mode("r") and not self._in_eof:
            diff = self._buffer.tell() - self._in_pos
            if diff > 0:
                junk = self._fileobj.read(diff)
                self._in_pos += len(junk)
                if len(junk) < diff:
                    self._in_eof = True
                    self._buffer.flush()

    def _seek(self, offset, whence):
        # Ensure we've read enough to simply do the seek on the buffer
        if self._check_mode("r") and not self._in_eof:
            if whence == 0:
                if offset > self._in_pos:
                    self._read_rest()
            if whence == 1:
                if self._buffer.tell() + offset > self._in_pos:
                    self._read_rest()
            if whence == 2:
                self._read_rest()
        # Then just do it on the buffer...
        self._buffer.seek(offset, whence)

    def _tell(self):
        return self._buffer.tell()

    def _truncate(self, size):
        if self._check_mode("r") and not self._in_eof:
            if size > self._in_pos:
                self._read_rest()
        self._in_eof = True
        try:
            self._buffer.truncate(size)
        except TypeError:
            et, ev, tb = sys.exc_info()
            # SpooledTemporaryFile.truncate() doesn't accept size paramter.
            try:
                self._buffer._file.truncate(size)
            except Exception:
                raise et, ev, tb
        # StringIO objects don't truncate to larger size correctly.
        if hasattr(self._buffer, "_file"):
            _file = self._buffer._file
            if hasattr(_file, "getvalue"):
                if len(_file.getvalue()) != size:
                    curpos = _file.tell()
                    _file.seek(0, 2)
                    _file.write("\x00" * (size - len(_file.getvalue())))
                    _file.seek(curpos)
        self._was_truncated = True

    def _read_rest(self):
        """Read the rest of the input stream."""
        if self._in_eof:
            return
        pos = self._buffer.tell()
        self._buffer.seek(0, 2)
        data = self._fileobj.read(self._bufsize)
        while data:
            self._in_pos += len(data)
            self._buffer.write(data)
            data = self._fileobj.read(self._bufsize)
        self._in_eof = True
        self._buffer.flush()
        self._buffer.seek(pos)