util.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153
  1. # -*- coding: utf-8 -*-
  2. # Natural Language Toolkit: Twitter client
  3. #
  4. # Copyright (C) 2001-2019 NLTK Project
  5. # Author: Ewan Klein <ewan@inf.ed.ac.uk>
  6. # Lorenzo Rubio <lrnzcig@gmail.com>
  7. # URL: <http://nltk.org/>
  8. # For license information, see LICENSE.TXT
  9. """
  10. Authentication utilities to accompany :module:`twitterclient`.
  11. """
  12. from __future__ import print_function
  13. import os
  14. import pprint
  15. from twython import Twython
  16. def credsfromfile(creds_file=None, subdir=None, verbose=False):
  17. """
  18. Convenience function for authentication
  19. """
  20. return Authenticate().load_creds(
  21. creds_file=creds_file, subdir=subdir, verbose=verbose
  22. )
  23. class Authenticate(object):
  24. """
  25. Methods for authenticating with Twitter.
  26. """
  27. def __init__(self):
  28. self.creds_file = 'credentials.txt'
  29. self.creds_fullpath = None
  30. self.oauth = {}
  31. try:
  32. self.twitter_dir = os.environ['TWITTER']
  33. self.creds_subdir = self.twitter_dir
  34. except KeyError:
  35. self.twitter_dir = None
  36. self.creds_subdir = None
  37. def load_creds(self, creds_file=None, subdir=None, verbose=False):
  38. """
  39. Read OAuth credentials from a text file.
  40. ::
  41. File format for OAuth 1
  42. =======================
  43. app_key=YOUR_APP_KEY
  44. app_secret=YOUR_APP_SECRET
  45. oauth_token=OAUTH_TOKEN
  46. oauth_token_secret=OAUTH_TOKEN_SECRET
  47. ::
  48. File format for OAuth 2
  49. =======================
  50. app_key=YOUR_APP_KEY
  51. app_secret=YOUR_APP_SECRET
  52. access_token=ACCESS_TOKEN
  53. :param str file_name: File containing credentials. ``None`` (default) reads\
  54. data from `TWITTER/'credentials.txt'`
  55. """
  56. if creds_file is not None:
  57. self.creds_file = creds_file
  58. if subdir is None:
  59. if self.creds_subdir is None:
  60. msg = (
  61. "Supply a value to the 'subdir' parameter or"
  62. + " set the TWITTER environment variable."
  63. )
  64. raise ValueError(msg)
  65. else:
  66. self.creds_subdir = subdir
  67. self.creds_fullpath = os.path.normpath(
  68. os.path.join(self.creds_subdir, self.creds_file)
  69. )
  70. if not os.path.isfile(self.creds_fullpath):
  71. raise OSError('Cannot find file {}'.format(self.creds_fullpath))
  72. with open(self.creds_fullpath) as infile:
  73. if verbose:
  74. print('Reading credentials file {}'.format(self.creds_fullpath))
  75. for line in infile:
  76. if '=' in line:
  77. name, value = line.split('=', 1)
  78. self.oauth[name.strip()] = value.strip()
  79. self._validate_creds_file(verbose=verbose)
  80. return self.oauth
  81. def _validate_creds_file(self, verbose=False):
  82. """Check validity of a credentials file."""
  83. oauth1 = False
  84. oauth1_keys = ['app_key', 'app_secret', 'oauth_token', 'oauth_token_secret']
  85. oauth2 = False
  86. oauth2_keys = ['app_key', 'app_secret', 'access_token']
  87. if all(k in self.oauth for k in oauth1_keys):
  88. oauth1 = True
  89. elif all(k in self.oauth for k in oauth2_keys):
  90. oauth2 = True
  91. if not (oauth1 or oauth2):
  92. msg = 'Missing or incorrect entries in {}\n'.format(self.creds_file)
  93. msg += pprint.pformat(self.oauth)
  94. raise ValueError(msg)
  95. elif verbose:
  96. print('Credentials file "{}" looks good'.format(self.creds_file))
  97. def add_access_token(creds_file=None):
  98. """
  99. For OAuth 2, retrieve an access token for an app and append it to a
  100. credentials file.
  101. """
  102. if creds_file is None:
  103. path = os.path.dirname(__file__)
  104. creds_file = os.path.join(path, 'credentials2.txt')
  105. oauth2 = credsfromfile(creds_file=creds_file)
  106. app_key = oauth2['app_key']
  107. app_secret = oauth2['app_secret']
  108. twitter = Twython(app_key, app_secret, oauth_version=2)
  109. access_token = twitter.obtain_access_token()
  110. tok = 'access_token={}\n'.format(access_token)
  111. with open(creds_file, 'a') as infile:
  112. print(tok, file=infile)
  113. def guess_path(pth):
  114. """
  115. If the path is not absolute, guess that it is a subdirectory of the
  116. user's home directory.
  117. :param str pth: The pathname of the directory where files of tweets should be written
  118. """
  119. if os.path.isabs(pth):
  120. return pth
  121. else:
  122. return os.path.expanduser(os.path.join("~", pth))