Add ability to persist base urls with clientid/secret/token (fixes #200)

3 files changed, 65 insertions, 19 deletions

diff --git a/mastodon/Mastodon.py b/mastodon/Mastodon.py
index 465beb7..c84ac6a 100644
--- a/mastodon/Mastodon.py
+++ b/mastodon/Mastodon.py
@@ -267,16 +267,17 @@ class Mastodon:
 
         if to_file is not None:
             with open(to_file, 'w') as secret_file:
-                secret_file.write(response['client_id'] + '\n')
-                secret_file.write(response['client_secret'] + '\n')
-
+                secret_file.write(response['client_id'] + "\n")
+                secret_file.write(response['client_secret'] + "\n")
+                secret_file.write(api_base_url + "\n")
+                
         return (response['client_id'], response['client_secret'])
 
     ###
     # Authentication, including constructor
     ###
     def __init__(self, client_id=None, client_secret=None, access_token=None,
-                 api_base_url=__DEFAULT_BASE_URL, debug_requests=False,
+                 api_base_url=None, debug_requests=False,
                  ratelimit_method="wait", ratelimit_pacefactor=1.1,
                  request_timeout=__DEFAULT_TIMEOUT, mastodon_version=None,
                  version_check_mode = "created", session=None):
@@ -285,9 +286,12 @@ class Mastodon:
         give a `client_id` and it is not a file, you must also give a secret. If you specify an
         `access_token` then you don't need to specify a `client_id`. It is allowed to specify
         neither - in this case, you will be restricted to only using endpoints that do not
-        require authentication.
+        require authentication.  If a file is given as `client_id`, client ID, secret and 
+        base url are read from that file.
 
-        You can also specify an `access_token`, directly or as a file (as written by `log_in()`_).
+        You can also specify an `access_token`, directly or as a file (as written by `log_in()`_). If
+        a file is given, Mastodon.py also tries to load the base URL from this file, if present. A
+        client id and secret are not required in this case.
 
         Mastodon.py can try to respect rate limits in several ways, controlled by `ratelimit_method`.
         "throw" makes functions throw a `MastodonRatelimitError` when the rate
@@ -298,8 +302,9 @@ class Mastodon:
         even in "wait" and "pace" mode, requests can still fail due to network or other problems! Also
         note that "pace" and "wait" are NOT thread safe.
 
-        Specify `api_base_url` if you wish to talk to an instance other than the flagship one.
-        If a file is given as `client_id`, client ID and secret are read from that file.
+        Specify `api_base_url` if you wish to talk to an instance other than the flagship one. When
+        reading from client id or access token files as written by Mastodon.py 1.5.0 or larger,
+        this can be omitted.
 
         By default, a timeout of 300 seconds is used for all requests. If you wish to change this,
         pass the desired timeout (in seconds) as `request_timeout`.
@@ -317,7 +322,10 @@ class Mastodon:
         changed after the version of Mastodon that is connected has been released. If it is set to "none",
         version checking is disabled.
         """
-        self.api_base_url = Mastodon.__protocolize(api_base_url)
+        self.api_base_url = None
+        if not api_base_url is None:
+            self.api_base_url = Mastodon.__protocolize(api_base_url)
+        
         self.client_id = client_id
         self.client_secret = client_secret
         self.access_token = access_token
@@ -364,6 +372,13 @@ class Mastodon:
                 with open(self.client_id, 'r') as secret_file:
                     self.client_id = secret_file.readline().rstrip()
                     self.client_secret = secret_file.readline().rstrip()
+                    
+                    try_base_url = secret_file.readline().rstrip()
+                    if (not try_base_url is None) and len(try_base_url) != 0:
+                        try_base_url = Mastodon.__protocolize(try_base_url)
+                        if not (self.api_base_url is None or try_base_url == self.api_base_url):
+                            raise MastodonIllegalArgumentError('Mismatch in base URLs between files and/or specified')
+                        self.api_base_url = try_base_url
             else:
                 if self.client_secret is None:
                     raise MastodonIllegalArgumentError('Specified client id directly, but did not supply secret')
@@ -371,7 +386,14 @@ class Mastodon:
         if self.access_token is not None and os.path.isfile(self.access_token):
             with open(self.access_token, 'r') as token_file:
                 self.access_token = token_file.readline().rstrip()
-    
+                
+                try_base_url = token_file.readline().rstrip()
+                if (not try_base_url is None) and len(try_base_url) != 0:
+                    try_base_url = Mastodon.__protocolize(try_base_url)
+                    if not (self.api_base_url is None or try_base_url == self.api_base_url):
+                         raise MastodonIllegalArgumentError('Mismatch in base URLs between files and/or specified')
+                    self.api_base_url = try_base_url
+                        
     def retrieve_mastodon_version(self):
         """
         Determine installed mastodon version and set major, minor and patch (not including RC info) accordingly.
@@ -508,8 +530,9 @@ class Mastodon:
 
         if to_file is not None:
             with open(to_file, 'w') as token_file:
-                token_file.write(response['access_token'] + '\n')
-        
+                token_file.write(response['access_token'] + "\n")
+                token_file.write(self.api_base_url + "\n")
+                
         self.__logged_in_id = None
         
         return response['access_token']
@@ -572,8 +595,9 @@ class Mastodon:
         
         if to_file is not None:
             with open(to_file, 'w') as token_file:
-                token_file.write(response['access_token'] + '\n')
-        
+                token_file.write(response['access_token'] + "\n")
+                token_file.write(self.api_base_url + "\n")
+                
         self.__logged_in_id = None
         
         return response['access_token']
diff --git a/tests/test_auth.py b/tests/test_auth.py
index b4a004e..fbf8974 100644
--- a/tests/test_auth.py
+++ b/tests/test_auth.py
@@ -30,7 +30,6 @@ def test_log_in_password(api_anonymous):
         password='mastodonadmin')
     assert token
 
-
 @pytest.mark.vcr()
 def test_log_in_password_incorrect(api_anonymous):
     with pytest.raises(MastodonIllegalArgumentError):
@@ -38,7 +37,6 @@ def test_log_in_password_incorrect(api_anonymous):
             username='admin@localhost',
             password='hunter2')
 
-
 @pytest.mark.vcr()
 def test_log_in_password_to_file(api_anonymous, tmpdir):
     filepath = tmpdir.join('token')
@@ -46,18 +44,42 @@ def test_log_in_password_to_file(api_anonymous, tmpdir):
         username='admin@localhost',
         password='mastodonadmin',
         to_file=str(filepath))
-    token = filepath.read_text('UTF-8').rstrip()
+    token = filepath.read_text('UTF-8').rstrip().split("\n")[0]
     assert token
     api = api_anonymous
     api.access_token = token
     assert api.account_verify_credentials()
 
+@pytest.mark.vcr()
+def test_url_errors(tmpdir):
+    clientid_good = tmpdir.join("clientid")
+    token_good = tmpdir.join("token")
+    clientid_bad = tmpdir.join("clientid_bad")
+    token_bad = tmpdir.join("token_bad")
+    
+    clientid_good.write_text("foo\nbar\nhttps://zombo.com\n", "UTF-8")
+    token_good.write_text("foo\nhttps://zombo.com\n", "UTF-8")
+    clientid_bad.write_text("foo\nbar\nhttps://evil.org\n", "UTF-8")
+    token_bad.write_text("foo\nhttps://evil.org\n", "UTF-8")  
+    
+    api = Mastodon(client_id = clientid_good, access_token = token_good)
+    assert api
+    assert api.api_base_url == "https://zombo.com"
+    assert Mastodon(client_id = clientid_good, access_token = token_good, api_base_url = "zombo.com")
+    
+    with pytest.raises(MastodonIllegalArgumentError):
+        Mastodon(client_id = clientid_good, access_token = token_bad, api_base_url = "zombo.com")
+        
+    with pytest.raises(MastodonIllegalArgumentError):
+        Mastodon(client_id = clientid_bad, access_token = token_good, api_base_url = "zombo.com")
+        
+    with pytest.raises(MastodonIllegalArgumentError):
+        Mastodon(client_id = clientid_bad, access_token = token_bad, api_base_url = "zombo.com")        
 
 @pytest.mark.skip(reason="Not sure how to test this without setting up selenium or a similar browser automation suite to click on the allow button")
 def test_log_in_code(api_anonymous):
     pass
 
-
 @pytest.mark.skip(reason="Not supported by Mastodon >:@ (yet?)")
 def test_log_in_refresh(api_anonymous):
     pass
diff --git a/tests/test_create_app.py b/tests/test_create_app.py
index b9de298..8260751 100644
--- a/tests/test_create_app.py
+++ b/tests/test_create_app.py
@@ -31,7 +31,7 @@ def test_create_app(mocker, to_file=None, redirect_uris=None, website=None):
 def test_create_app_to_file(mocker, tmpdir):
     filepath = tmpdir.join('credentials')
     test_create_app(mocker, to_file=str(filepath))
-    assert filepath.read_text('UTF-8') == "foo\nbar\n"
+    assert filepath.read_text('UTF-8') == "foo\nbar\nhttps://example.com\n"
 
 def test_create_app_redirect_uris(mocker):
     test_create_app(mocker, redirect_uris='http://example.net')