package net.lamgc.scext.onedrive_transfer import com.microsoft.aad.msal4j.ClientCredentialFactory import com.microsoft.aad.msal4j.ConfidentialClientApplication import com.microsoft.aad.msal4j.IAccount import com.microsoft.aad.msal4j.SilentParameters import com.microsoft.graph.authentication.IAuthenticationProvider import com.microsoft.graph.httpcore.HttpClients import com.microsoft.graph.models.Drive import com.microsoft.graph.requests.GraphServiceClient import mu.KotlinLogging import okhttp3.Request import org.jetbrains.exposed.sql.Database import org.jetbrains.exposed.sql.SchemaUtils import org.jetbrains.exposed.sql.transactions.transaction import org.telegram.abilitybots.api.bot.BaseAbilityBot import org.telegram.telegrambots.meta.api.objects.Document import java.net.URL import java.util.concurrent.CompletableFuture class OneDriveTransferService( private val bot: BaseAbilityBot, private val config: ExtensionConfig, private val db: Database ) { private val logger = KotlinLogging.logger { } val accountManager: OneDriveTransferSettingManager private val authClient: ConfidentialClientApplication = ConfidentialClientApplication.builder( config.clientId, ClientCredentialFactory.createFromSecret(config.clientSecret), ) .authority(OneDriveTransferSettingManager.AUTHORITY) .setTokenCacheAccessAspect(DatabaseTokenCache(db)) .build() init { transaction(db) { SchemaUtils.create(OneDriveTransferSettings, TokenCaches) } accountManager = OneDriveTransferSettingManager(authClient, db) } fun createGraphClient(userId: Long): GraphServiceClient { val cache = THREAD_CURRENT_GRAPH_CLIENT.get() if (cache?.tgUserId == userId) { return cache.client } val serviceClient = accountManager.getTransferSetting(userId)?.let { authClient.accounts.get().firstOrNull { account -> account.homeAccountId() == it.accountId } }?.let { Companion.createGraphClient(authClient, it) } ?: throw OneDriveNotLoginException() THREAD_CURRENT_GRAPH_CLIENT.set(ClientCache(userId, serviceClient)) return serviceClient } fun createLoginUrl(userId: Long) = accountManager.createAuthorizationRequest(userId) fun updateAccount(userId: Long, redirectUrl: URL) = accountManager.updateAccount(userId, redirectUrl) fun listDriversByUserId(userId: Long): List { val sites = try { val sites = createGraphClient(userId).sites() .buildRequest().get()?.currentPage ?: emptyList() sites.map { it.drive }.filterNotNull().toList() } catch (e: Exception) { logger.debug(e) { "获取 OneDrive 站点失败, 可能是用户没有权限或不是组织账号." } emptyList() } val drive = createGraphClient(userId).me().drive().buildRequest().get() val drives = createGraphClient(userId).drives() .buildRequest() .get()?.currentPage ?: emptyList() return mutableListOf().apply { if (drive != null) { add(drive) } addAll(drives) addAll(sites) } } fun setDrive(userId: Long, driveId: String) { val transferSetting = accountManager.getTransferSetting(userId) ?: throw OneDriveNotLoginException() accountManager.doSomething { transferSetting.driveId = driveId } } fun getCurrentDrive(userId: Long): Drive? { val transferSetting = accountManager.getTransferSetting(userId) ?: throw OneDriveNotLoginException() val graphClient = createGraphClient(userId) return graphClient.drives(transferSetting.driveId) .buildRequest() .get() } fun submitUploadDocumentTask(userId: Long, messageId: Int, document: Document) { val transferSetting = accountManager.getTransferSetting(userId) ?: throw OneDriveNotLoginException() OneDriveTransferCenter.submitUploadTask( OneDriveTransferTask( userId, bot, this, document, transferSetting.driveId, transferSetting.storagePath, ).apply { extra["chatId"] = userId extra["messageId"] = messageId } ) } companion object { private val THREAD_CURRENT_GRAPH_CLIENT = ThreadLocal() fun createGraphClient(authClient: ConfidentialClientApplication, iAccount: IAccount): GraphServiceClient { return GraphServiceClient.builder() .httpClient(HttpClients.createDefault(MsalAuthorizationProvider(authClient, iAccount))) .buildClient() } } } private data class ClientCache( val tgUserId: Long, val client: GraphServiceClient, ) class MsalAuthorizationProvider(private val authClientApplication: ConfidentialClientApplication, private val iAccount: IAccount) : IAuthenticationProvider { override fun getAuthorizationTokenAsync(requestUrl: URL): CompletableFuture { return authClientApplication.acquireTokenSilently( SilentParameters.builder(OneDriveTransferSettingManager.OAUTH2_SCOPE, iAccount).build() ).thenApply { it.accessToken() } } }