1# Copyright (c) 2013 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
5"""Handle gdata spreadsheet authentication."""
6
7import BaseHTTPServer
8import httplib
9import os
10import Queue
11import threading
12import socket
13import time
14
15import gdata.gauth
16import gdata.spreadsheets.client
17
18# Local port of http daemon for receiving the refresh token from the redirected
19# url returned by authentication server.
20START_PORT = 12345
21ACCESS_CODE_Q = Queue.Queue()
22
23# Token storage
24TOKEN_STORAGE_NAME = '%s/.spreadsheets.oauth2.dat' % os.getenv('HOME')
25
26# Application's CLIENT_ID and SECRET for OAuthToken which is created from
27# google's api console's API access - https://code.google.com/apis/console
28CLIENT_ID = '657833351030.apps.googleusercontent.com'
29CLIENT_SECRET = 'h72FzPdzfbN3I4U3M3l1DSiT'
30
31USER_AGENT = 'Pressure Calibration Data Collector'
32SCOPES = 'https://spreadsheets.google.com/feeds/ https://docs.google.com/feeds/'
33RETRY = 10
34
35
36class AuthenticationHandler(BaseHTTPServer.BaseHTTPRequestHandler):
37    """Authentication handler class."""
38
39    def do_QUIT(self):
40        """Do QUIT reuqest."""
41        self.send_response(200)
42        self.end_headers()
43        self.server.stop = True
44
45    def do_GET(self):  # pylint: disable=g-bad-name
46        """Do GET request."""
47        self.send_response(200)
48        self.send_header('Content-type', 'text/html')
49        self.end_headers()
50        self.port = START_PORT
51        query_str = self.path[2:]
52        if query_str.startswith('code='):
53            self.wfile.write('Spreadsheet authentication complete, '
54                             'please back to command prompt.')
55            ACCESS_CODE_Q.put(query_str[5:])
56            return
57        if query_str.startswith('error='):
58            print "Ouch, error: '%s'." % query_str[6:]
59            raise Exception("Exception during approval process: '%s'." %
60                            query_str[6:])
61
62    def log_message(self, format, *args):  # pylint: disable=redefined-builtin
63        pass
64
65
66class AuthenticationHTTPD(BaseHTTPServer.HTTPServer):
67    """ Handle redirected response from authentication server."""
68    def serve_forever(self, poll_interval=0.5):
69        poll_interval = poll_interval
70        self.stop = False
71        while not self.stop:
72            self.handle_request()
73
74
75class AuthenticationServer(threading.Thread):
76    """Authentication http server thread."""
77
78    def __init__(self):
79        self.authserver = None
80        self.started = False
81        threading.Thread.__init__(self)
82
83    def run(self):
84        for ports in range(START_PORT, START_PORT + RETRY):
85            self.port = ports
86            try:
87                self.authserver = AuthenticationHTTPD(('', self.port),
88                                                      AuthenticationHandler)
89                self.started = True
90                self.authserver.serve_forever()
91                return
92
93            except socket.error, se:
94                # port is busy, there must be another instance running...
95                if self.port == START_PORT + RETRY - 1:
96                    raise se
97                else:
98                    continue  # keep trying new ports
99
100            except KeyboardInterrupt:
101                print '^C received, shutting down authentication server.'
102                self.stop = True
103                return  # out of retry loop
104
105            except Exception:
106                self.stop = True
107                return  # out of retry loop
108
109    def get_port(self):
110        """Get the running port number."""
111        for retry in xrange(RETRY):
112            if self.started and self.authserver:
113                return self.port
114            else:
115                time.sleep(retry * 2)
116                continue
117
118
119class SpreadsheetAuthorizer:
120    """ Handle gdata api authentication for spreadsheet client."""
121    def __init__(self):
122        self.lock = threading.Lock()
123        self.refresh_token = None
124        self.redirect_url = None
125        self.httpd_auth = None
126        self.port = None
127
128    def _start_server(self):
129        """Start http daemon for handling refresh token."""
130        if not self.httpd_auth or not self.httpd_auth.isAlive():
131            ### Starting webserver if necessary
132            self.httpd_auth = AuthenticationServer()
133            self.httpd_auth.start()
134            self.port = self.httpd_auth.get_port()
135
136    def _stop_server(self):
137        """Stop http daemon."""
138        if self.httpd_auth:
139            try:
140                conn = httplib.HTTPConnection('localhost:%s' % self.port)
141                conn.request('QUIT', '/')
142                conn.getresponse()
143                time.sleep(1)
144                del self.httpd_auth
145                self.httpd_auth = None
146            except Exception, e:
147                print "Failed to quit local auth server...'%s'." % e
148        return
149
150    def authorize(self, ss_client):
151        """Authorize the spreadsheet client.
152
153        @param ss_client: spreadsheet client
154        """
155        self._read_refresh_token()
156        token = gdata.gauth.OAuth2Token(CLIENT_ID,
157                                        CLIENT_SECRET,
158                                        SCOPES,
159                                        USER_AGENT,
160                                        refresh_token = self.refresh_token)
161        try:
162            if not self.refresh_token:
163                self._start_server()
164                token = gdata.gauth.OAuth2Token(CLIENT_ID,
165                                                CLIENT_SECRET,
166                                                SCOPES,
167                                                USER_AGENT)
168                redirect_url = 'http://localhost:' + str(self.port)
169                url = token.generate_authorize_url(redirect_url)
170                print ('\nPlease open the following URL and use @chromium.org'
171                       'account for authentication and authorization of the'
172                       'spreadsheet access:\n' + url)
173                print 'Waiting for you to authenticate...'
174                while ACCESS_CODE_Q.empty():
175                    time.sleep(.25)
176                access_code = ACCESS_CODE_Q.get()
177                print 'ACCESS CODE is ' + access_code
178                token.get_access_token(access_code)
179                self.refresh_token = token.refresh_token
180                print 'REFRESH TOKEN is ' + self.refresh_token
181                self._write_refresh_token()
182                self._stop_server()
183            token.authorize(ss_client)
184            return True
185        except IOError:
186            print "ERROR!!!!!!!!!!!!!!!!"
187        return False
188
189    def _read_refresh_token(self):
190        """Read refresh token from storage."""
191        try:
192            self.lock.acquire()
193            token_file = open(TOKEN_STORAGE_NAME, 'r')
194            self.refresh_token = token_file.readline().strip()
195            token_file.close()
196            self.lock.release()
197            return self.refresh_token
198        except IOError:
199            self.lock.release()
200            return None
201
202    def _write_refresh_token(self):
203        """Write refresh token into storage."""
204        try:
205            self.lock.acquire()
206            token_descriptor = os.open(TOKEN_STORAGE_NAME,
207                                       os.O_CREAT | os.O_WRONLY | os.O_TRUNC,
208                                       0600)
209            token_file = os.fdopen(token_descriptor, 'w')
210            token_file.write(self.refresh_token + '\n')
211            token_file.close()
212            self.lock.release()
213        except (IOError, OSError):
214            self.lock.release()
215            print 'Error, can not write refresh token\n'
216