Repository: Backblaze/B2_Command_Line_Tool

Issue:
Automatic re-authorizing with application key does not work
When it re-authorizes, it should use the same application key used for the original authentication, but it's using the account ID.

Discussion:

Patch:
diff --git a/b2/account_info/abstract.py b/b2/account_info/abstract.py
index 3b1f77f..44bbdbd 100644
--- a/b2/account_info/abstract.py
+++ b/b2/account_info/abstract.py
@@ -88,6 +88,10 @@ class AbstractAccountInfo(object):
     def get_account_id(self):
         """ returns account_id or raises MissingAccountData exception """
 
+    @abstractmethod
+    def get_account_id_or_app_key_id(self):
+        """ returns the account id or key id used to authenticate """
+
     @abstractmethod
     def get_account_auth_token(self):
         """ returns account_auth_token or raises MissingAccountData exception """
@@ -133,6 +137,7 @@ class AbstractAccountInfo(object):
         application_key,
         realm,
         allowed=None,
+        account_id_or_app_key_id=None,
     ):
         """
         Stores the results of b2_authorize_account.
@@ -157,6 +162,7 @@ class AbstractAccountInfo(object):
             application_key,
             realm,
             allowed,
+            account_id_or_app_key_id,
         )
 
     @classmethod
@@ -176,8 +182,16 @@ class AbstractAccountInfo(object):
 
     @abstractmethod
     def _set_auth_data(
-        self, account_id, auth_token, api_url, download_url, minimum_part_size, application_key,
-        realm, allowed
+        self,
+        account_id,
+        auth_token,
+        api_url,
+        download_url,
+        minimum_part_size,
+        application_key,
+        realm,
+        allowed,
+        account_id_or_app_key_id,
     ):
         """
         Stores the auth data.  Can assume that 'allowed' is present and valid.
diff --git a/b2/account_info/in_memory.py b/b2/account_info/in_memory.py
index 2289ad3..045f56b 100644
--- a/b2/account_info/in_memory.py
+++ b/b2/account_info/in_memory.py
@@ -38,6 +38,7 @@ class InMemoryAccountInfo(UrlPoolAccountInfo):
 
     def _clear_in_memory_account_fields(self):
         self._account_id = None
+        self._account_id_or_app_key_id = None
         self._allowed = None
         self._api_url = None
         self._application_key = None
@@ -57,8 +58,10 @@ class InMemoryAccountInfo(UrlPoolAccountInfo):
         application_key,
         realm,
         allowed,
+        account_id_or_app_key_id,
     ):
         self._account_id = account_id
+        self._account_id_or_app_key_id = account_id_or_app_key_id
         self._auth_token = auth_token
         self._api_url = api_url
         self._download_url = download_url
@@ -84,6 +87,10 @@ class InMemoryAccountInfo(UrlPoolAccountInfo):
     def get_account_id(self):
         return self._account_id
 
+    @_raise_missing_if_result_is_none
+    def get_account_id_or_app_key_id(self):
+        return self._account_id_or_app_key_id
+
     @_raise_missing_if_result_is_none
     def get_account_auth_token(self):
         return self._auth_token
diff --git a/b2/account_info/sqlite_account_info.py b/b2/account_info/sqlite_account_info.py
index 72733f3..3e9ddaf 100644
--- a/b2/account_info/sqlite_account_info.py
+++ b/b2/account_info/sqlite_account_info.py
@@ -37,7 +37,11 @@ class SqliteAccountInfo(UrlPoolAccountInfo):
     completed.
     """
 
-    def __init__(self, file_name=None):
+    def __init__(self, file_name=None, last_upgrade_to_run=None):
+        """
+        :param file_name: The sqlite file to use; overrides the default.
+        :param last_upgrade_to_run: For testing only, override the auto-update on the db.
+        """
         self.thread_local = threading.local()
         user_account_info_path = file_name or os.environ.get(
             B2_ACCOUNT_INFO_ENV_VAR, B2_ACCOUNT_INFO_DEFAULT_FILE
@@ -45,23 +49,23 @@ class SqliteAccountInfo(UrlPoolAccountInfo):
         self.filename = file_name or os.path.expanduser(user_account_info_path)
         self._validate_database()
         with self._get_connection() as conn:
-            self._create_tables(conn)
+            self._create_tables(conn, last_upgrade_to_run)
         super(SqliteAccountInfo, self).__init__()
 
-    def _validate_database(self):
+    def _validate_database(self, last_upgrade_to_run=None):
         """
         Makes sure that the database is openable.  Removes the file if it's not.
         """
         # If there is no file there, that's fine.  It will get created when
         # we connect.
         if not os.path.exists(self.filename):
-            self._create_database()
+            self._create_database(last_upgrade_to_run)
             return
 
         # If we can connect to the database, and do anything, then all is good.
         try:
             with self._connect() as conn:
-                self._create_tables(conn)
+                self._create_tables(conn, last_upgrade_to_run)
                 return
         except sqlite3.DatabaseError:
             pass  # fall through to next case
@@ -79,10 +83,10 @@ class SqliteAccountInfo(UrlPoolAccountInfo):
                     # remove the json file
                     os.unlink(self.filename)
                     # create a database
-                    self._create_database()
+                    self._create_database(last_upgrade_to_run)
                     # add the data from the JSON file
                     with self._connect() as conn:
-                        self._create_tables(conn)
+                        self._create_tables(conn, last_upgrade_to_run)
                         insert_statement = """
                             INSERT INTO account
                             (account_id, application_key, account_auth_token, api_url, download_url, minimum_part_size, realm)
@@ -111,7 +115,7 @@ class SqliteAccountInfo(UrlPoolAccountInfo):
     def _connect(self):
         return sqlite3.connect(self.filename, isolation_level='EXCLUSIVE')
 
-    def _create_database(self):
+    def _create_database(self, last_upgrade_to_run):
         """
         Makes sure that the database is created and sets the file permissions.
         This should be done before storing any sensitive data in it.
@@ -120,14 +124,14 @@ class SqliteAccountInfo(UrlPoolAccountInfo):
         conn = self._connect()
         try:
             with conn:
-                self._create_tables(conn)
+                self._create_tables(conn, last_upgrade_to_run)
         finally:
             conn.close()
 
         # Set the file permissions
         os.chmod(self.filename, stat.S_IRUSR | stat.S_IWUSR)
 
-    def _create_tables(self, conn):
+    def _create_tables(self, conn, last_upgrade_to_run):
         conn.execute(
             """
             CREATE TABLE IF NOT EXISTS
@@ -172,8 +176,14 @@ class SqliteAccountInfo(UrlPoolAccountInfo):
            );
         """
         )
+        # By default, we run all the upgrades
+        last_upgrade_to_run = 2 if last_upgrade_to_run is None else last_upgrade_to_run
         # Add the 'allowed' column if it hasn't been yet.
-        self._ensure_update(1, 'ALTER TABLE account ADD COLUMN allowed TEXT;')
+        if 1 <= last_upgrade_to_run:
+            self._ensure_update(1, 'ALTER TABLE account ADD COLUMN allowed TEXT;')
+        # Add the 'account_id_or_app_key_id' column if it hasn't been yet
+        if 2 <= last_upgrade_to_run:
+            self._ensure_update(2, 'ALTER TABLE account ADD COLUMN account_id_or_app_key_id TEXT;')
 
     def _ensure_update(self, update_number, update_command):
         """
@@ -212,6 +222,7 @@ class SqliteAccountInfo(UrlPoolAccountInfo):
         application_key,
         realm,
         allowed,
+        account_id_or_app_key_id,
     ):
         assert self.allowed_is_valid(allowed)
         with self._get_connection() as conn:
@@ -220,14 +231,53 @@ class SqliteAccountInfo(UrlPoolAccountInfo):
             conn.execute('DELETE FROM bucket_upload_url;')
             insert_statement = """
                 INSERT INTO account
-                (account_id, application_key, account_auth_token, api_url, download_url, minimum_part_size, realm, allowed)
-                values (?, ?, ?, ?, ?, ?, ?, ?);
+                (account_id, account_id_or_app_key_id, application_key, account_auth_token, api_url, download_url, minimum_part_size, realm, allowed)
+                values (?, ?, ?, ?, ?, ?, ?, ?, ?);
+            """
+
+            conn.execute(
+                insert_statement, (
+                    account_id,
+                    account_id_or_app_key_id,
+                    application_key,
+                    auth_token,
+                    api_url,
+                    download_url,
+                    minimum_part_size,
+                    realm,
+                    json.dumps(allowed),
+                )
+            )
+
+    def set_auth_data_with_schema_0_for_test(
+        self,
+        account_id,
+        auth_token,
+        api_url,
+        download_url,
+        minimum_part_size,
+        application_key,
+        realm,
+    ):
+        with self._get_connection() as conn:
+            conn.execute('DELETE FROM account;')
+            conn.execute('DELETE FROM bucket;')
+            conn.execute('DELETE FROM bucket_upload_url;')
+            insert_statement = """
+                INSERT INTO account
+                (account_id, application_key, account_auth_token, api_url, download_url, minimum_part_size, realm)
+                values (?, ?, ?, ?, ?, ?, ?);
             """
 
             conn.execute(
                 insert_statement, (
-                    account_id, application_key, auth_token, api_url, download_url,
-                    minimum_part_size, realm, json.dumps(allowed)
+                    account_id,
+                    application_key,
+                    auth_token,
+                    api_url,
+                    download_url,
+                    minimum_part_size,
+                    realm,
                 )
             )
 
@@ -237,6 +287,16 @@ class SqliteAccountInfo(UrlPoolAccountInfo):
     def get_account_id(self):
         return self._get_account_info_or_raise('account_id')
 
+    def get_account_id_or_app_key_id(self):
+        """
+        The 'account_id_or_app_key_id' column was not in the original schema, so it may be NULL.
+        """
+        result = self._get_account_info_or_raise('account_id_or_app_key_id')
+        if result is None:
+            return self.get_account_id()
+        else:
+            return result
+
     def get_api_url(self):
         return self._get_account_info_or_raise('api_url')
 
diff --git a/b2/api.py b/b2/api.py
index a1400e1..644191e 100644
--- a/b2/api.py
+++ b/b2/api.py
@@ -103,7 +103,7 @@ class B2Api(object):
         try:
             self.authorize_account(
                 self.account_info.get_realm(),
-                self.account_info.get_account_id(),
+                self.account_info.get_account_id_or_app_key_id(),
                 self.account_info.get_application_key(),
             )
         except MissingAccountData:
@@ -163,6 +163,7 @@ class B2Api(object):
             application_key,
             realm,
             allowed,
+            account_id_or_key_id,
         )
 
     def get_account_id(self):
diff --git a/b2/raw_simulator.py b/b2/raw_simulator.py
index 0fcb999..7513b23 100644
--- a/b2/raw_simulator.py
+++ b/b2/raw_simulator.py
@@ -542,6 +542,9 @@ class RawSimulator(AbstractRawApi):
         # Map from auth token to the KeySimulator for it.
         self.auth_token_to_key = dict()
 
+        # Set of auth tokens that have expired
+        self.expired_auth_tokens = set()
+
         # Counter for generating auth tokens.
         self.auth_token_counter = 0
 
@@ -556,6 +559,16 @@ class RawSimulator(AbstractRawApi):
         self.app_key_counter = 0
         self.upload_errors = []
 
+    def expire_auth_token(self, auth_token):
+        """
+        Simulate the auth token expiring.
+
+        The next call that tries to use this auth token will get an
+        auth_token_expired error.
+        """
+        assert auth_token in self.auth_token_to_key
+        self.expired_auth_tokens.add(auth_token)
+
     def create_account(self):
         """
         Returns (accountId, masterApplicationKey) for a newly created account.
@@ -930,6 +943,8 @@ class RawSimulator(AbstractRawApi):
         assert key_sim is not None
         assert api_url == self.API_URL
         assert account_id == key_sim.account_id
+        if account_auth_token in self.expired_auth_tokens:
+            raise InvalidAuthToken('auth token expired', 'auth_token_expired')
         if capability not in key_sim.capabilities:
             raise Unauthorized('', 'unauthorized')
         if key_sim.bucket_id_or_none is not None and key_sim.bucket_id_or_none != bucket_id:


Test patch:
diff --git a/test/stub_account_info.py b/test/stub_account_info.py
index e8b2347..0d469b2 100644
--- a/test/stub_account_info.py
+++ b/test/stub_account_info.py
@@ -34,6 +34,7 @@ class StubAccountInfo(AbstractAccountInfo):
         self.download_url = None
         self.minimum_part_size = None
         self.realm = None
+        self.account_id_or_app_key_id = None
         self._large_file_uploads = collections.defaultdict(list)
         self._large_file_uploads_lock = threading.Lock()
 
@@ -51,6 +52,7 @@ class StubAccountInfo(AbstractAccountInfo):
         application_key,
         realm,
         allowed,
+        account_id_or_app_key_id,
     ):
         self.account_id = account_id
         self.auth_token = auth_token
@@ -60,6 +62,7 @@ class StubAccountInfo(AbstractAccountInfo):
         self.application_key = application_key
         self.realm = realm
         self.allowed = allowed
+        self.account_id_or_app_key_id = account_id_or_app_key_id
 
     def refresh_entire_bucket_name_cache(self, name_id_iterable):
         self.buckets = {}
@@ -82,6 +85,9 @@ class StubAccountInfo(AbstractAccountInfo):
     def get_account_id(self):
         return self.account_id
 
+    def get_account_id_or_app_key_id(self):
+        return self.account_id_or_app_key_id
+
     def get_account_auth_token(self):
         return self.auth_token
 
diff --git a/test/test_account_info.py b/test/test_account_info.py
index 923269c..d0fa8f7 100644
--- a/test/test_account_info.py
+++ b/test/test_account_info.py
@@ -261,11 +261,48 @@ class TestSqliteAccountInfo(AccountInfoBase, TestBase):
         account_info = self._make_info()
         self.assertEqual('auth_token', account_info.get_account_auth_token())
 
+    def test_upgrade_1_default_allowed(self):
+        """
+        The 'allowed' field should be the default for upgraded databases.
+        """
+        old_account_info = self._make_sqlite_account_info(last_upgrade_to_run=0)
+        old_account_info.set_auth_data_with_schema_0_for_test(
+            'account_id',
+            'auth_token',
+            'api_url',
+            'dowload_url',
+            100,  # minimum part size
+            'application_key',
+            'realm',
+        )
+        new_account_info = self._make_info()
+        self.assertEqual(AbstractAccountInfo.DEFAULT_ALLOWED, new_account_info.get_allowed())
+
+    def test_upgrade_2_default_app_key(self):
+        """
+        The 'account_id_or_app_key_id' field should default to the account id.
+        """
+        old_account_info = self._make_sqlite_account_info(last_upgrade_to_run=0)
+        old_account_info.set_auth_data_with_schema_0_for_test(
+            'account_id',
+            'auth_token',
+            'api_url',
+            'dowload_url',
+            100,  # minimum part size
+            'application_key',
+            'realm',
+        )
+        new_account_info = self._make_info()
+        self.assertEqual('account_id', new_account_info.get_account_id_or_app_key_id())
+
     def _make_info(self):
+        return self._make_sqlite_account_info()
+
+    def _make_sqlite_account_info(self, last_upgrade_to_run=None):
         """
         Returns a new StoredAccountInfo that has just read the data from the file.
         """
-        return SqliteAccountInfo(file_name=self.db_path)
+        return SqliteAccountInfo(file_name=self.db_path, last_upgrade_to_run=None)
 
     def test_account_info_persistence(self):
         self._test_account_info(check_persistence=True)
diff --git a/test/test_api.py b/test/test_api.py
index f72c336..fb2a016 100644
--- a/test/test_api.py
+++ b/test/test_api.py
@@ -42,6 +42,20 @@ class TestApi(TestBase):
             [b.name for b in self.api.list_buckets(bucket_name='bucket1')],
         )
 
+    def test_reauthorize_with_app_key(self):
+        # authorize and create a key
+        self._authorize_account()
+        key = self.api.create_key(['listBuckets'], 'key1')
+
+        # authorize with the key
+        self.api.authorize_account('production', key['applicationKeyId'], key['applicationKey'])
+
+        # expire the auth token we just got
+        self.raw_api.expire_auth_token(self.account_info.get_account_auth_token())
+
+        # listing buckets should work, after it re-authorizes
+        self.api.list_buckets()
+
     def test_list_buckets_with_restriction(self):
         self._authorize_account()
         bucket1 = self.api.create_bucket('bucket1', 'allPrivate')