dotfiles

My personal dotfiles
git clone git://pcardensab.com/dotfiles
Log | Files | Refs | Submodules | README

mutt_oauth2.py (17888B)


      1 #!/usr/bin/env python3
      2 #
      3 # Mutt OAuth2 token management script, version 2020-08-07
      4 # Written against python 3.7.3, not tried with earlier python versions.
      5 #
      6 #   Copyright (C) 2020 Alexander Perlis
      7 #
      8 #   This program is free software; you can redistribute it and/or
      9 #   modify it under the terms of the GNU General Public License as
     10 #   published by the Free Software Foundation; either version 2 of the
     11 #   License, or (at your option) any later version.
     12 #
     13 #   This program is distributed in the hope that it will be useful,
     14 #   but WITHOUT ANY WARRANTY; without even the implied warranty of
     15 #   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
     16 #   General Public License for more details.
     17 #
     18 #   You should have received a copy of the GNU General Public License
     19 #   along with this program; if not, write to the Free Software
     20 #   Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
     21 #   02110-1301, USA.
     22 '''Mutt OAuth2 token management'''
     23 
     24 import sys
     25 import json
     26 import argparse
     27 import urllib.parse
     28 import urllib.request
     29 import imaplib
     30 import poplib
     31 import smtplib
     32 import base64
     33 import secrets
     34 import hashlib
     35 import time
     36 from datetime import timedelta, datetime
     37 from pathlib import Path
     38 import socket
     39 import http.server
     40 import subprocess
     41 
     42 # The token file must be encrypted because it contains multi-use bearer tokens
     43 # whose usage does not require additional verification. Specify whichever
     44 # encryption and decryption pipes you prefer. They should read from standard
     45 # input and write to standard output. The example values here invoke GPG,
     46 # although won't work until an appropriate identity appears in the first line.
     47 ENCRYPTION_PIPE = ['gpg', '--encrypt', '--recipient', '99E05F373B9E3FCB78CA93EEAB3EA9F44C8DD92B']
     48 DECRYPTION_PIPE = ['gpg', '--decrypt']
     49 
     50 google_client = json.load(open(Path.home() / '.config/mutt/clients/google.json'))
     51 
     52 registrations = {
     53     'google': {
     54         'authorize_endpoint': 'https://accounts.google.com/o/oauth2/auth',
     55         'devicecode_endpoint': 'https://oauth2.googleapis.com/device/code',
     56         'token_endpoint': 'https://accounts.google.com/o/oauth2/token',
     57         'redirect_uri': 'urn:ietf:wg:oauth:2.0:oob',
     58         'imap_endpoint': 'imap.gmail.com',
     59         'pop_endpoint': 'pop.gmail.com',
     60         'smtp_endpoint': 'smtp.gmail.com',
     61         'sasl_method': 'OAUTHBEARER',
     62         'scope': 'https://mail.google.com/',
     63         'client_id': google_client['installed']['client_id'],
     64         'client_secret': google_client['installed']['client_secret'],
     65     },
     66     'microsoft': {
     67         'authorize_endpoint': 'https://login.microsoftonline.com/common/oauth2/v2.0/authorize',
     68         'devicecode_endpoint': 'https://login.microsoftonline.com/common/oauth2/v2.0/devicecode',
     69         'token_endpoint': 'https://login.microsoftonline.com/common/oauth2/v2.0/token',
     70         'redirect_uri': 'https://login.microsoftonline.com/common/oauth2/nativeclient',
     71         'tenant': 'common',
     72         'imap_endpoint': 'outlook.office365.com',
     73         'pop_endpoint': 'outlook.office365.com',
     74         'smtp_endpoint': 'smtp.office365.com',
     75         'sasl_method': 'XOAUTH2',
     76         'scope': ('offline_access https://outlook.office.com/IMAP.AccessAsUser.All '
     77                   'https://outlook.office.com/POP.AccessAsUser.All '
     78                   'https://outlook.office.com/SMTP.Send'),
     79         'client_id': '',
     80         'client_secret': '',
     81     },
     82 }
     83 
     84 ap = argparse.ArgumentParser(epilog='''
     85 This script obtains and prints a valid OAuth2 access token.  State is maintained in an
     86 encrypted TOKENFILE.  Run with "--verbose --authorize" to get started or whenever all
     87 tokens have expired, optionally with "--authflow" to override the default authorization
     88 flow.  To truly start over from scratch, first delete TOKENFILE.  Use "--verbose --test"
     89 to test the IMAP/POP/SMTP endpoints.
     90 ''')
     91 ap.add_argument('-v', '--verbose', action='store_true', help='increase verbosity')
     92 ap.add_argument('-d', '--debug', action='store_true', help='enable debug output')
     93 ap.add_argument('tokenfile', help='persistent token storage')
     94 ap.add_argument('-a', '--authorize', action='store_true', help='manually authorize new tokens')
     95 ap.add_argument('--authflow', help='authcode | localhostauthcode | devicecode')
     96 ap.add_argument('-t', '--test', action='store_true', help='test IMAP/POP/SMTP endpoints')
     97 args = ap.parse_args()
     98 
     99 token = {}
    100 path = Path(args.tokenfile)
    101 if path.exists():
    102     if 0o777 & path.stat().st_mode != 0o600:
    103         sys.exit('Token file has unsafe mode. Suggest deleting and starting over.')
    104     try:
    105         sub = subprocess.run(DECRYPTION_PIPE, check=True, input=path.read_bytes(),
    106                              capture_output=True)
    107         token = json.loads(sub.stdout)
    108     except subprocess.CalledProcessError:
    109         sys.exit('Difficulty decrypting token file. Is your decryption agent primed for '
    110                  'non-interactive usage, or an appropriate environment variable such as '
    111                  'GPG_TTY set to allow interactive agent usage from inside a pipe?')
    112 
    113 
    114 def writetokenfile():
    115     '''Writes global token dictionary into token file.'''
    116     if not path.exists():
    117         path.touch(mode=0o600)
    118     if 0o777 & path.stat().st_mode != 0o600:
    119         sys.exit('Token file has unsafe mode. Suggest deleting and starting over.')
    120     sub2 = subprocess.run(ENCRYPTION_PIPE, check=True, input=json.dumps(token).encode(),
    121                           capture_output=True)
    122     path.write_bytes(sub2.stdout)
    123 
    124 
    125 if args.debug:
    126     print('Obtained from token file:', json.dumps(token))
    127 if not token:
    128     if not args.authorize:
    129         sys.exit('You must run script with "--authorize" at least once.')
    130     print('Available app and endpoint registrations:', *registrations)
    131     token['registration'] = input('OAuth2 registration: ')
    132     token['authflow'] = input('Preferred OAuth2 flow ("authcode" or "localhostauthcode" '
    133                               'or "devicecode"): ')
    134     token['email'] = input('Account e-mail address: ')
    135     token['access_token'] = ''
    136     token['access_token_expiration'] = ''
    137     token['refresh_token'] = ''
    138     writetokenfile()
    139 
    140 if token['registration'] not in registrations:
    141     sys.exit(f'ERROR: Unknown registration "{token["registration"]}". Delete token file '
    142              f'and start over.')
    143 registration = registrations[token['registration']]
    144 
    145 authflow = token['authflow']
    146 if args.authflow:
    147     authflow = args.authflow
    148 
    149 baseparams = {'client_id': registration['client_id']}
    150 # Microsoft uses 'tenant' but Google does not
    151 if 'tenant' in registration:
    152     baseparams['tenant'] = registration['tenant']
    153 
    154 
    155 def access_token_valid():
    156     '''Returns True when stored access token exists and is still valid at this time.'''
    157     token_exp = token['access_token_expiration']
    158     return token_exp and datetime.now() < datetime.fromisoformat(token_exp)
    159 
    160 
    161 def update_tokens(r):
    162     '''Takes a response dictionary, extracts tokens out of it, and updates token file.'''
    163     token['access_token'] = r['access_token']
    164     token['access_token_expiration'] = (datetime.now() +
    165                                         timedelta(seconds=int(r['expires_in']))).isoformat()
    166     if 'refresh_token' in r:
    167         token['refresh_token'] = r['refresh_token']
    168     writetokenfile()
    169     if args.verbose:
    170         print(f'NOTICE: Obtained new access token, expires {token["access_token_expiration"]}.')
    171 
    172 
    173 if args.authorize:
    174     p = baseparams.copy()
    175     p['scope'] = registration['scope']
    176 
    177     if authflow in ('authcode', 'localhostauthcode'):
    178         verifier = secrets.token_urlsafe(90)
    179         challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest())[:-1]
    180         redirect_uri = registration['redirect_uri']
    181         listen_port = 0
    182         if authflow == 'localhostauthcode':
    183             # Find an available port to listen on
    184             s = socket.socket()
    185             s.bind(('127.0.0.1', 0))
    186             listen_port = s.getsockname()[1]
    187             s.close()
    188             redirect_uri = 'http://localhost:'+str(listen_port)+'/'
    189             # Probably should edit the port number into the actual redirect URL.
    190 
    191         p.update({'login_hint': token['email'],
    192                   'response_type': 'code',
    193                   'redirect_uri': redirect_uri,
    194                   'code_challenge': challenge,
    195                   'code_challenge_method': 'S256'})
    196         print(registration["authorize_endpoint"] + '?' +
    197               urllib.parse.urlencode(p, quote_via=urllib.parse.quote))
    198 
    199         authcode = ''
    200         if authflow == 'authcode':
    201             authcode = input('Visit displayed URL to retrieve authorization code. Enter '
    202                              'code from server (might be in browser address bar): ')
    203         else:
    204             print('Visit displayed URL to authorize this application. Waiting...',
    205                   end='', flush=True)
    206 
    207             class MyHandler(http.server.BaseHTTPRequestHandler):
    208                 '''Handles the browser query resulting from redirect to redirect_uri.'''
    209 
    210                 # pylint: disable=C0103
    211                 def do_HEAD(self):
    212                     '''Response to a HEAD requests.'''
    213                     self.send_response(200)
    214                     self.send_header('Content-type', 'text/html')
    215                     self.end_headers()
    216 
    217                 def do_GET(self):
    218                     '''For GET request, extract code parameter from URL.'''
    219                     # pylint: disable=W0603
    220                     global authcode
    221                     querystring = urllib.parse.urlparse(self.path).query
    222                     querydict = urllib.parse.parse_qs(querystring)
    223                     if 'code' in querydict:
    224                         authcode = querydict['code'][0]
    225                     self.do_HEAD()
    226                     self.wfile.write(b'<html><head><title>Authorizaton result</title></head>')
    227                     self.wfile.write(b'<body><p>Authorization redirect completed. You may '
    228                                      b'close this window.</p></body></html>')
    229             with http.server.HTTPServer(('127.0.0.1', listen_port), MyHandler) as httpd:
    230                 try:
    231                     httpd.handle_request()
    232                 except KeyboardInterrupt:
    233                     pass
    234 
    235         if not authcode:
    236             sys.exit('Did not obtain an authcode.')
    237 
    238         for k in 'response_type', 'login_hint', 'code_challenge', 'code_challenge_method':
    239             del p[k]
    240         p.update({'grant_type': 'authorization_code',
    241                   'code': authcode,
    242                   'client_secret': registration['client_secret'],
    243                   'code_verifier': verifier})
    244         try:
    245             response = urllib.request.urlopen(registration['token_endpoint'],
    246                                               urllib.parse.urlencode(p).encode())
    247         except urllib.error.HTTPError as err:
    248             print(err.code, err.reason)
    249             response = err
    250         response = response.read()
    251         if args.debug:
    252             print(response)
    253         response = json.loads(response)
    254         if 'error' in response:
    255             print(response['error'])
    256             if 'error_description' in response:
    257                 print(response['error_description'])
    258             sys.exit(1)
    259 
    260     elif authflow == 'devicecode':
    261         try:
    262             response = urllib.request.urlopen(registration['devicecode_endpoint'],
    263                                               urllib.parse.urlencode(p).encode())
    264         except urllib.error.HTTPError as err:
    265             print(err.code, err.reason)
    266             response = err
    267         response = response.read()
    268         if args.debug:
    269             print(response)
    270         response = json.loads(response)
    271         if 'error' in response:
    272             print(response['error'])
    273             if 'error_description' in response:
    274                 print(response['error_description'])
    275             sys.exit(1)
    276         print(response['message'])
    277         del p['scope']
    278         p.update({'grant_type': 'urn:ietf:params:oauth:grant-type:device_code',
    279                   'client_secret': registration['client_secret'],
    280                   'device_code': response['device_code']})
    281         interval = int(response['interval'])
    282         print('Polling...', end='', flush=True)
    283         while True:
    284             time.sleep(interval)
    285             print('.', end='', flush=True)
    286             try:
    287                 response = urllib.request.urlopen(registration['token_endpoint'],
    288                                                   urllib.parse.urlencode(p).encode())
    289             except urllib.error.HTTPError as err:
    290                 # Not actually always an error, might just mean "keep trying..."
    291                 response = err
    292             response = response.read()
    293             if args.debug:
    294                 print(response)
    295             response = json.loads(response)
    296             if 'error' not in response:
    297                 break
    298             if response['error'] == 'authorization_declined':
    299                 print(' user declined authorization.')
    300                 sys.exit(1)
    301             if response['error'] == 'expired_token':
    302                 print(' too much time has elapsed.')
    303                 sys.exit(1)
    304             if response['error'] != 'authorization_pending':
    305                 print(response['error'])
    306                 if 'error_description' in response:
    307                     print(response['error_description'])
    308                 sys.exit(1)
    309         print()
    310 
    311     else:
    312         sys.exit(f'ERROR: Unknown OAuth2 flow "{token["authflow"]}. Delete token file and '
    313                  f'start over.')
    314 
    315     update_tokens(response)
    316 
    317 
    318 if not access_token_valid():
    319     if args.verbose:
    320         print('NOTICE: Invalid or expired access token; using refresh token '
    321               'to obtain new access token.')
    322     if not token['refresh_token']:
    323         sys.exit('ERROR: No refresh token. Run script with "--authorize".')
    324     p = baseparams.copy()
    325     p.update({'client_secret': registration['client_secret'],
    326               'refresh_token': token['refresh_token'],
    327               'grant_type': 'refresh_token'})
    328     try:
    329         response = urllib.request.urlopen(registration['token_endpoint'],
    330                                           urllib.parse.urlencode(p).encode())
    331     except urllib.error.HTTPError as err:
    332         print(err.code, err.reason)
    333         response = err
    334     response = response.read()
    335     if args.debug:
    336         print(response)
    337     response = json.loads(response)
    338     if 'error' in response:
    339         print(response['error'])
    340         if 'error_description' in response:
    341             print(response['error_description'])
    342         print('Perhaps refresh token invalid. Try running once with "--authorize"')
    343         sys.exit(1)
    344     update_tokens(response)
    345 
    346 
    347 if not access_token_valid():
    348     sys.exit('ERROR: No valid access token. This should not be able to happen.')
    349 
    350 
    351 if args.verbose:
    352     print('Access Token: ', end='')
    353 print(token['access_token'])
    354 
    355 
    356 def build_sasl_string(user, host, port, bearer_token):
    357     '''Build appropriate SASL string, which depends on cloud server's supported SASL method.'''
    358     if registration['sasl_method'] == 'OAUTHBEARER':
    359         return f'n,a={user},\1host={host}\1port={port}\1auth=Bearer {bearer_token}\1\1'
    360     if registration['sasl_method'] == 'XOAUTH2':
    361         return f'user={user}\1auth=Bearer {bearer_token}\1\1'
    362     sys.exit(f'Unknown SASL method {registration["sasl_method"]}.')
    363 
    364 
    365 if args.test:
    366     errors = False
    367 
    368     imap_conn = imaplib.IMAP4_SSL(registration['imap_endpoint'])
    369     sasl_string = build_sasl_string(token['email'], registration['imap_endpoint'], 993,
    370                                     token['access_token'])
    371     if args.debug:
    372         imap_conn.debug = 4
    373     try:
    374         imap_conn.authenticate(registration['sasl_method'], lambda _: sasl_string.encode())
    375         # Microsoft has a bug wherein a mismatch between username and token can still report a
    376         # successful login... (Try a consumer login with the token from a work/school account.)
    377         # Fortunately subsequent commands fail with an error. Thus we follow AUTH with another
    378         # IMAP command before reporting success.
    379         imap_conn.list()
    380         if args.verbose:
    381             print('IMAP authentication succeeded')
    382     except imaplib.IMAP4.error as e:
    383         print('IMAP authentication FAILED (does your account allow IMAP?):', e)
    384         errors = True
    385 
    386     pop_conn = poplib.POP3_SSL(registration['pop_endpoint'])
    387     sasl_string = build_sasl_string(token['email'], registration['pop_endpoint'], 995,
    388                                     token['access_token'])
    389     if args.debug:
    390         pop_conn.set_debuglevel(2)
    391     try:
    392         # poplib doesn't have an auth command taking an authenticator object
    393         # Microsoft requires a two-line SASL for POP
    394         # pylint: disable=W0212
    395         pop_conn._shortcmd('AUTH ' + registration['sasl_method'])
    396         pop_conn._shortcmd(base64.standard_b64encode(sasl_string.encode()).decode())
    397         if args.verbose:
    398             print('POP authentication succeeded')
    399     except poplib.error_proto as e:
    400         print('POP authentication FAILED (does your account allow POP?):', e.args[0].decode())
    401         errors = True
    402 
    403     # SMTP_SSL would be simpler but Microsoft does not answer on port 465.
    404     smtp_conn = smtplib.SMTP(registration['smtp_endpoint'], 587)
    405     sasl_string = build_sasl_string(token['email'], registration['smtp_endpoint'], 587,
    406                                     token['access_token'])
    407     smtp_conn.ehlo('test')
    408     smtp_conn.starttls()
    409     smtp_conn.ehlo('test')
    410     if args.debug:
    411         smtp_conn.set_debuglevel(2)
    412     try:
    413         smtp_conn.auth(registration['sasl_method'], lambda _=None: sasl_string)
    414         if args.verbose:
    415             print('SMTP authentication succeeded')
    416     except smtplib.SMTPAuthenticationError as e:
    417         print('SMTP authentication FAILED:', e)
    418         errors = True
    419 
    420     if errors:
    421         sys.exit(1)