1# Copyright 2018 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import socket
6import threading
7import Queue
8
9_BUF_SIZE = 4096
10
11class FakePrinter():
12    """
13    A fake printer (server).
14
15    It starts a thread that listens on given localhost's port and saves
16    incoming documents in the internal queue. Documents can be fetched from
17    the queue by calling the fetch_document() method. At the end, the printer
18    must be stopped by calling the stop() method. The stop() method is called
19    automatically when the object is managed by "with" statement.
20    See test_fake_printer.py for examples.
21
22    """
23
24    def __init__(self, port):
25        """
26        Initialize fake printer.
27
28        It configures the socket and starts the printer. If no exceptions
29        are thrown (the method succeeded), the printer must be stopped by
30        calling the stop() method.
31
32        @param port: port number on which the printer is supposed to listen
33
34        @raises socket or thread related exception in case of failure
35
36        """
37        # If set to True, the printer is stopped either by invoking stop()
38        # method or by an internal error
39        self._stopped = False
40        # It is set when printer is stopped because of some internal error
41        self._error_message = None
42        # An internal queue with printed documents
43        self._documents = Queue.Queue()
44        # Create a TCP/IP socket
45        self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
46        try:
47            # Bind the socket to the port
48            self._socket.bind( ('localhost', port) )
49            # Start thread
50            self._thread = threading.Thread(target = self._thread_read_docs)
51            self._thread.start();
52        except:
53            # failure - the socket must be closed before exit
54            self._socket.close()
55            raise
56
57
58    # These methods allow to use the 'with' statement to automaticaly stop
59    # the printer
60    def __enter__(self):
61        return self
62    def __exit__(self, exc_type, exc_value, traceback):
63        self.stop()
64
65
66    def stop(self):
67        """
68        Stops the printer.
69
70        """
71        self._stopped = True
72        self._thread.join()
73
74
75    def fetch_document(self, timeout):
76        """
77        Fetches the next document from the internal queue.
78
79        This method returns the next document and removes it from the internal
80        queue. If there is no documents in the queue, it blocks until one
81        arrives. If waiting time exceeds a given timeout, an exception is
82        raised.
83
84        @param timeout: max waiting time in seconds
85
86        @returns next document from the internal queue
87
88        @raises Exception if the timeout was reached
89
90        """
91        try:
92            return self._documents.get(block=True, timeout=timeout)
93        except Queue.Empty:
94            # Builds a message for the exception
95            message = 'Timeout occured when waiting for the document. '
96            if self._stopped:
97                message += 'The fake printer was stopped '
98                if self._error_message is None:
99                    message += 'by the stop() method.'
100                else:
101                    message += 'because of the error: %s.' % self._error_message
102            else:
103                message += 'The fake printer is in valid state.'
104            # Raises and exception
105            raise Exception(message)
106
107
108    def _read_whole_document(self):
109        """
110        Reads a document from the printer's socket.
111
112        It assumes that operation on sockets may timeout.
113
114        @returns whole document or None, if the printer was stopped
115
116        """
117        # Accepts incoming connection
118        while True:
119            try:
120                (connection, client_address) = self._socket.accept()
121                # success - exit the loop
122                break
123            except socket.timeout:
124                # exit if the printer was stopped, else return to the loop
125                if self._stopped:
126                    return None
127
128        # Reads document
129        document = ''
130        while True:
131            try:
132                data = connection.recv(_BUF_SIZE)
133                # success - check data and continue
134                if not data:
135                    # we got the whole document - exit the loop
136                    break
137                # save chunk of the document and return to the loop
138                document += data
139            except socket.timeout:
140                # exit if the printer was stopped, else return to the loop
141                if self._stopped:
142                    connection.close()
143                    return None
144
145        # Closes connection & returns document
146        connection.close()
147        return document
148
149
150    def _thread_read_docs(self):
151        """
152        Reads documents from the printer's socket and adds them to the
153        internal queue.
154
155        It exits when the printer is stopped by the stop() method.
156        In case of any error (exception) it stops the printer and exits.
157
158        """
159        try:
160            # Listen for incoming printer request.
161            self._socket.listen(1)
162            # All following socket's methods throw socket.timeout after
163            # 500 miliseconds
164            self._socket.settimeout(0.5)
165
166            while True:
167                # Reads document from the socket
168                document = self._read_whole_document()
169                # 'None' means that the printer was stopped -> exit
170                if document is None:
171                    break
172                # Adds documents to the internal queue
173                self._documents.put(document)
174        except BaseException as e:
175            # Error occured, the printer must be stopped -> exit
176            self._error_message = str(e)
177            self._stopped = True
178
179        # Closes socket before the exit
180        self._socket.close()
181