Updating server to handle lock_off with Hydrus-Key auth
This commit is contained in:
parent
a04a5c271c
commit
7067a5a7ac
|
@ -1,6 +1,7 @@
|
|||
import collections
|
||||
from . import HydrusConstants as HC
|
||||
from . import HydrusExceptions
|
||||
import hashlib
|
||||
import os
|
||||
import queue
|
||||
import re
|
||||
|
@ -26,20 +27,55 @@ class HydrusSessionManagerServer( object ):
|
|||
HG.controller.sub( self, 'RefreshAllAccounts', 'update_all_session_accounts' )
|
||||
|
||||
|
||||
def _GetAccountFromAccountKey( self, service_key, account_key ):
|
||||
|
||||
account_keys_to_accounts = self._service_keys_to_account_keys_to_accounts[ service_key ]
|
||||
|
||||
if account_key not in account_keys_to_accounts:
|
||||
|
||||
if HG.server_busy.locked():
|
||||
|
||||
raise HydrusExceptions.ServerBusyException( 'Sorry, server is busy and cannot fetch account data right now!' )
|
||||
|
||||
|
||||
account = HG.controller.Read( 'account', service_key, account_key )
|
||||
|
||||
account_keys_to_accounts[ account_key ] = account
|
||||
|
||||
|
||||
account = account_keys_to_accounts[ account_key ]
|
||||
|
||||
return account
|
||||
|
||||
|
||||
def _GetAccountKeyFromAccessKey( self, service_key, access_key ):
|
||||
|
||||
hashed_access_key = hashlib.sha256( access_key ).digest()
|
||||
|
||||
if hashed_access_key not in self._service_keys_to_hashed_access_keys_to_account_keys[ service_key ]:
|
||||
|
||||
if HG.server_busy.locked():
|
||||
|
||||
raise HydrusExceptions.ServerBusyException( 'Sorry, server is busy and cannot fetch account key data right now!' )
|
||||
|
||||
|
||||
account_key = HG.controller.Read( 'account_key_from_access_key', service_key, access_key )
|
||||
|
||||
self._service_keys_to_hashed_access_keys_to_account_keys[ service_key ][ hashed_access_key ] = account_key
|
||||
|
||||
|
||||
account_key = self._service_keys_to_hashed_access_keys_to_account_keys[ service_key ][ hashed_access_key ]
|
||||
|
||||
return account_key
|
||||
|
||||
|
||||
def AddSession( self, service_key, access_key ):
|
||||
|
||||
with self._lock:
|
||||
|
||||
account_key = HG.controller.Read( 'account_key_from_access_key', service_key, access_key )
|
||||
account_key = self._GetAccountKeyFromAccessKey( service_key, access_key )
|
||||
|
||||
account_keys_to_accounts = self._service_keys_to_account_keys_to_accounts[ service_key ]
|
||||
|
||||
if account_key not in account_keys_to_accounts:
|
||||
|
||||
account = HG.controller.Read( 'account', service_key, account_key )
|
||||
|
||||
account_keys_to_accounts[ account_key ] = account
|
||||
|
||||
account = self._GetAccountFromAccountKey( service_key, account_key )
|
||||
|
||||
session_key = HydrusData.GenerateKey()
|
||||
|
||||
|
@ -85,18 +121,9 @@ class HydrusSessionManagerServer( object ):
|
|||
|
||||
with self._lock:
|
||||
|
||||
account_key = HG.controller.Read( 'account_key_from_access_key', service_key, access_key )
|
||||
account_key = self._GetAccountKeyFromAccessKey( service_key, access_key )
|
||||
|
||||
account_keys_to_accounts = self._service_keys_to_account_keys_to_accounts[ service_key ]
|
||||
|
||||
if account_key not in account_keys_to_accounts:
|
||||
|
||||
account = HG.controller.Read( 'account', service_key, account_key )
|
||||
|
||||
account_keys_to_accounts[ account_key ] = account
|
||||
|
||||
|
||||
account = account_keys_to_accounts[ account_key ]
|
||||
account = self._GetAccountFromAccountKey( service_key, account_key )
|
||||
|
||||
return account
|
||||
|
||||
|
@ -108,9 +135,9 @@ class HydrusSessionManagerServer( object ):
|
|||
|
||||
service_keys_to_dirty_accounts = {}
|
||||
|
||||
for ( service_key, account_keys_to_accounts ) in list(self._service_keys_to_account_keys_to_accounts.items()):
|
||||
for ( service_key, account_keys_to_accounts ) in self._service_keys_to_account_keys_to_accounts.items():
|
||||
|
||||
dirty_accounts = [ account_key for account_key in list(account_keys_to_accounts.values()) if account_key.IsDirty() ]
|
||||
dirty_accounts = [ account_key for account_key in account_keys_to_accounts.values() if account_key.IsDirty() ]
|
||||
|
||||
if len( dirty_accounts ) > 0:
|
||||
|
||||
|
@ -130,7 +157,7 @@ class HydrusSessionManagerServer( object ):
|
|||
|
||||
if account_keys is None:
|
||||
|
||||
account_keys = list(account_keys_to_accounts.keys())
|
||||
account_keys = list( account_keys_to_accounts.keys() )
|
||||
|
||||
|
||||
for account_key in account_keys:
|
||||
|
@ -152,6 +179,8 @@ class HydrusSessionManagerServer( object ):
|
|||
|
||||
self._service_keys_to_account_keys_to_accounts = collections.defaultdict( dict )
|
||||
|
||||
self._service_keys_to_hashed_access_keys_to_account_keys = collections.defaultdict( dict )
|
||||
|
||||
existing_sessions = HG.controller.Read( 'sessions' )
|
||||
|
||||
else:
|
||||
|
@ -160,16 +189,26 @@ class HydrusSessionManagerServer( object ):
|
|||
|
||||
del self._service_keys_to_account_keys_to_accounts[ service_key ]
|
||||
|
||||
del self._service_keys_to_hashed_access_keys_to_account_keys[ service_key ]
|
||||
|
||||
existing_sessions = HG.controller.Read( 'sessions', service_key )
|
||||
|
||||
|
||||
for ( session_key, service_key, account, expires ) in existing_sessions:
|
||||
for ( session_key, service_key, account, hashed_access_key, expires ) in existing_sessions:
|
||||
|
||||
account_key = account.GetAccountKey()
|
||||
|
||||
self._service_keys_to_session_keys_to_sessions[ service_key ][ session_key ] = ( account_key, expires )
|
||||
|
||||
self._service_keys_to_account_keys_to_accounts[ service_key ][ account_key ] = account
|
||||
if account_key not in self._service_keys_to_account_keys_to_accounts:
|
||||
|
||||
self._service_keys_to_account_keys_to_accounts[ service_key ][ account_key ] = account
|
||||
|
||||
|
||||
if hashed_access_key not in self._service_keys_to_hashed_access_keys_to_account_keys:
|
||||
|
||||
self._service_keys_to_hashed_access_keys_to_account_keys[ service_key ][ hashed_access_key ] = account_key
|
||||
|
||||
|
||||
|
||||
|
||||
|
|
|
@ -564,7 +564,7 @@ class DB( HydrusDB.HydrusDB ):
|
|||
|
||||
#
|
||||
|
||||
hashed_account_key = hashlib.sha256( access_key ).digest()
|
||||
hashed_access_key = hashlib.sha256( access_key ).digest()
|
||||
|
||||
account_type = self._GetAccountTypeFromCache( service_id, account_type_id )
|
||||
|
||||
|
@ -576,7 +576,7 @@ class DB( HydrusDB.HydrusDB ):
|
|||
|
||||
dictionary_string = dictionary.DumpToString()
|
||||
|
||||
self._c.execute( 'INSERT INTO accounts ( service_id, account_key, hashed_access_key, account_type_id, created, expires, dictionary_string ) VALUES ( ?, ?, ?, ?, ?, ?, ? );', ( service_id, sqlite3.Binary( account_key ), sqlite3.Binary( hashed_account_key ), account_type_id, created, expires, dictionary_string ) )
|
||||
self._c.execute( 'INSERT INTO accounts ( service_id, account_key, hashed_access_key, account_type_id, created, expires, dictionary_string ) VALUES ( ?, ?, ?, ?, ?, ?, ? );', ( service_id, sqlite3.Binary( account_key ), sqlite3.Binary( hashed_access_key ), account_type_id, created, expires, dictionary_string ) )
|
||||
|
||||
else:
|
||||
|
||||
|
@ -1010,6 +1010,8 @@ class DB( HydrusDB.HydrusDB ):
|
|||
|
||||
account_ids_to_accounts = {}
|
||||
|
||||
account_ids_to_hashed_access_keys = {}
|
||||
|
||||
for ( session_key, service_id, account_id, expires ) in results:
|
||||
|
||||
if service_id not in service_ids_to_service_keys:
|
||||
|
@ -1028,7 +1030,16 @@ class DB( HydrusDB.HydrusDB ):
|
|||
|
||||
account = account_ids_to_accounts[ account_id ]
|
||||
|
||||
sessions.append( ( session_key, service_key, account, expires ) )
|
||||
if account_id not in account_ids_to_hashed_access_keys:
|
||||
|
||||
( hashed_access_key, ) = self._c.execute( 'SELECT hashed_access_key FROM accounts WHERE account_id = ?;', ( account_id, ) ).fetchone()
|
||||
|
||||
account_ids_to_hashed_access_keys[ account_id ] = hashed_access_key
|
||||
|
||||
|
||||
hashed_access_key = account_ids_to_hashed_access_keys[ account_id ]
|
||||
|
||||
sessions.append( ( session_key, service_key, account, hashed_access_key, expires ) )
|
||||
|
||||
|
||||
return sessions
|
||||
|
|
|
@ -1,5 +1,6 @@
|
|||
from . import ClientConstants as CC
|
||||
import collections
|
||||
import hashlib
|
||||
from . import HydrusConstants as HC
|
||||
from . import HydrusExceptions
|
||||
from . import HydrusNetwork
|
||||
|
@ -20,17 +21,27 @@ class TestSessions( unittest.TestCase ):
|
|||
|
||||
permissions = [ HC.GET_DATA, HC.POST_DATA, HC.POST_PETITIONS, HC.RESOLVE_PETITIONS, HC.MANAGE_USERS, HC.GENERAL_ADMIN, HC.EDIT_SERVICES ]
|
||||
|
||||
access_key = HydrusData.GenerateKey()
|
||||
account_key = HydrusData.GenerateKey()
|
||||
account_type = HydrusNetwork.AccountType.GenerateAdminAccountType( HC.SERVER_ADMIN )
|
||||
created = HydrusData.GetNow() - 100000
|
||||
expires = HydrusData.GetNow() + 300
|
||||
|
||||
account = HydrusNetwork.Account( account_key, account_type, created, expires )
|
||||
account_key_1 = HydrusData.GenerateKey()
|
||||
account_key_2 = HydrusData.GenerateKey()
|
||||
|
||||
access_key_1 = HydrusData.GenerateKey()
|
||||
hashed_access_key_1 = hashlib.sha256( access_key_1 ).digest()
|
||||
|
||||
access_key_2 = HydrusData.GenerateKey()
|
||||
hashed_access_key_2 = hashlib.sha256( access_key_2 ).digest()
|
||||
|
||||
account = HydrusNetwork.Account( account_key_1, account_type, created, expires )
|
||||
account_2 = HydrusNetwork.Account( account_key_2, account_type, created, expires )
|
||||
|
||||
# test timeout
|
||||
|
||||
expires = HydrusData.GetNow() - 10
|
||||
|
||||
HG.test_controller.SetRead( 'sessions', [ ( session_key_1, service_key, account, expires ) ] )
|
||||
HG.test_controller.SetRead( 'sessions', [ ( session_key_1, service_key, account, hashed_access_key_1, expires ) ] )
|
||||
|
||||
session_manager = HydrusSessions.HydrusSessionManagerServer()
|
||||
|
||||
|
@ -39,11 +50,18 @@ class TestSessions( unittest.TestCase ):
|
|||
session_manager.GetAccount( service_key, session_key_1 )
|
||||
|
||||
|
||||
# test missing
|
||||
|
||||
with self.assertRaises( HydrusExceptions.SessionException ):
|
||||
|
||||
session_manager.GetAccount( service_key, HydrusData.GenerateKey() )
|
||||
|
||||
|
||||
# test fetching a session already in db, after bootup
|
||||
|
||||
expires = HydrusData.GetNow() + 300
|
||||
|
||||
HG.test_controller.SetRead( 'sessions', [ ( session_key_1, service_key, account, expires ) ] )
|
||||
HG.test_controller.SetRead( 'sessions', [ ( session_key_1, service_key, account, hashed_access_key_1, expires ) ] )
|
||||
|
||||
session_manager = HydrusSessions.HydrusSessionManagerServer()
|
||||
|
||||
|
@ -51,20 +69,38 @@ class TestSessions( unittest.TestCase ):
|
|||
|
||||
self.assertIs( read_account, account )
|
||||
|
||||
read_account = session_manager.GetAccountFromAccessKey( service_key, access_key_1 )
|
||||
|
||||
self.assertIs( read_account, account )
|
||||
|
||||
# test too busy to add a new session for a new account it doesn't know about
|
||||
|
||||
HG.server_busy.acquire()
|
||||
|
||||
with self.assertRaises( HydrusExceptions.ServerBusyException ):
|
||||
|
||||
session_manager.AddSession( service_key, HydrusData.GenerateKey() )
|
||||
|
||||
session_manager.GetAccountFromAccessKey( service_key, HydrusData.GenerateKey() )
|
||||
|
||||
|
||||
# but ok to get for a session that already exists while busy
|
||||
|
||||
session_manager.GetAccount( service_key, session_key_1 )
|
||||
session_manager.GetAccountFromAccessKey( service_key, access_key_1 )
|
||||
|
||||
HG.server_busy.release()
|
||||
|
||||
# test adding a session
|
||||
|
||||
HG.test_controller.ClearWrites( 'session' )
|
||||
|
||||
expires = HydrusData.GetNow() + 300
|
||||
|
||||
account_key_2 = HydrusData.GenerateKey()
|
||||
|
||||
account_2 = HydrusNetwork.Account( account_key_2, account_type, created, expires )
|
||||
|
||||
HG.test_controller.SetRead( 'account_key_from_access_key', account_key_2 )
|
||||
HG.test_controller.SetRead( 'account', account_2 )
|
||||
|
||||
( session_key_2, expires_2 ) = session_manager.AddSession( service_key, access_key )
|
||||
( session_key_2, expires_2 ) = session_manager.AddSession( service_key, access_key_2 )
|
||||
|
||||
[ ( args, kwargs ) ] = HG.test_controller.GetWrite( 'session' )
|
||||
|
||||
|
@ -76,64 +112,84 @@ class TestSessions( unittest.TestCase ):
|
|||
|
||||
self.assertIs( read_account, account_2 )
|
||||
|
||||
read_account = session_manager.GetAccountFromAccessKey( service_key, access_key_2 )
|
||||
|
||||
self.assertIs( read_account, account_2 )
|
||||
|
||||
# test adding a new session for an account already in the manager
|
||||
|
||||
HG.test_controller.SetRead( 'account_key_from_access_key', account_key )
|
||||
HG.test_controller.SetRead( 'account_key_from_access_key', account_key_1 )
|
||||
HG.test_controller.SetRead( 'account', account )
|
||||
|
||||
( session_key_3, expires_3 ) = session_manager.AddSession( service_key, access_key )
|
||||
( session_key_3, expires_3 ) = session_manager.AddSession( service_key, access_key_1 )
|
||||
|
||||
[ ( args, kwargs ) ] = HG.test_controller.GetWrite( 'session' )
|
||||
|
||||
( written_session_key, written_service_key, written_account_key, written_expires ) = args
|
||||
|
||||
self.assertEqual( ( session_key_3, service_key, account_key, expires_3 ), ( written_session_key, written_service_key, written_account_key, written_expires ) )
|
||||
self.assertEqual( ( session_key_3, service_key, account_key_1, expires_3 ), ( written_session_key, written_service_key, written_account_key, written_expires ) )
|
||||
|
||||
read_account = session_manager.GetAccount( service_key, session_key_1 )
|
||||
|
||||
self.assertIs( read_account, account )
|
||||
|
||||
read_account = session_manager.GetAccount( service_key, session_key_3 )
|
||||
|
||||
self.assertIs( read_account, account )
|
||||
|
||||
read_account_original = session_manager.GetAccount( service_key, session_key_1 )
|
||||
read_account = session_manager.GetAccountFromAccessKey( service_key, access_key_1 )
|
||||
|
||||
self.assertIs( read_account, read_account_original )
|
||||
self.assertIs( read_account, account )
|
||||
|
||||
# test individual account refresh
|
||||
|
||||
expires = HydrusData.GetNow() + 300
|
||||
|
||||
updated_account = HydrusNetwork.Account( account_key, account_type, created, expires )
|
||||
new_obj_account_1 = HydrusNetwork.Account( account_key_1, account_type, created, expires )
|
||||
|
||||
HG.test_controller.SetRead( 'account', updated_account )
|
||||
HG.test_controller.SetRead( 'account', new_obj_account_1 )
|
||||
|
||||
session_manager.RefreshAccounts( service_key, [ account_key ] )
|
||||
session_manager.RefreshAccounts( service_key, [ account_key_1 ] )
|
||||
|
||||
read_account = session_manager.GetAccount( service_key, session_key_1 )
|
||||
|
||||
self.assertIs( read_account, updated_account )
|
||||
self.assertIs( read_account, new_obj_account_1 )
|
||||
|
||||
read_account = session_manager.GetAccount( service_key, session_key_3 )
|
||||
|
||||
self.assertIs( read_account, updated_account )
|
||||
self.assertIs( read_account, new_obj_account_1 )
|
||||
|
||||
read_account = session_manager.GetAccountFromAccessKey( service_key, access_key_1 )
|
||||
|
||||
self.assertIs( read_account, new_obj_account_1 )
|
||||
|
||||
# test all account refresh
|
||||
|
||||
expires = HydrusData.GetNow() + 300
|
||||
|
||||
updated_account_2 = HydrusNetwork.Account( account_key, account_type, created, expires )
|
||||
new_obj_account_2 = HydrusNetwork.Account( account_key_2, account_type, created, expires )
|
||||
|
||||
HG.test_controller.SetRead( 'sessions', [ ( session_key_1, service_key, updated_account_2, expires ), ( session_key_2, service_key, account_2, expires ), ( session_key_3, service_key, updated_account_2, expires ) ] )
|
||||
HG.test_controller.SetRead( 'sessions', [ ( session_key_1, service_key, new_obj_account_2, hashed_access_key_2, expires ), ( session_key_2, service_key, new_obj_account_1, hashed_access_key_1, expires ), ( session_key_3, service_key, new_obj_account_2, hashed_access_key_2, expires ) ] )
|
||||
|
||||
session_manager.RefreshAllAccounts()
|
||||
|
||||
read_account = session_manager.GetAccount( service_key, session_key_1 )
|
||||
|
||||
self.assertIs( read_account, updated_account_2 )
|
||||
self.assertIs( read_account, new_obj_account_2 )
|
||||
|
||||
read_account = session_manager.GetAccount( service_key, session_key_2 )
|
||||
|
||||
self.assertIs( read_account, account_2 )
|
||||
self.assertIs( read_account, new_obj_account_1 )
|
||||
|
||||
read_account = session_manager.GetAccount( service_key, session_key_3 )
|
||||
|
||||
self.assertIs( read_account, updated_account_2 )
|
||||
self.assertIs( read_account, new_obj_account_2 )
|
||||
|
||||
read_account = session_manager.GetAccountFromAccessKey( service_key, access_key_1 )
|
||||
|
||||
self.assertIs( read_account, new_obj_account_1 )
|
||||
|
||||
read_account = session_manager.GetAccountFromAccessKey( service_key, access_key_2 )
|
||||
|
||||
self.assertIs( read_account, new_obj_account_2 )
|
||||
|
||||
|
|
Loading…
Reference in New Issue