diff --git a/build.gradle.kts b/build.gradle.kts index afdf303..7017bf3 100644 --- a/build.gradle.kts +++ b/build.gradle.kts @@ -20,7 +20,6 @@ repositories { dependencies { compileOnly("org.slf4j:slf4j-api:2.0.10") compileOnly("io.github.microutils:kotlin-logging:3.0.5") - implementation("ch.qos.logback:logback-classic:1.4.14") compileOnly("net.lamgc:scalabot-extension:0.6.1") implementation("com.fasterxml.jackson.core:jackson-core:2.16.1") diff --git a/src/main/kotlin/Databases.kt b/src/main/kotlin/Databases.kt index d2a3dbf..b55f9da 100644 --- a/src/main/kotlin/Databases.kt +++ b/src/main/kotlin/Databases.kt @@ -2,6 +2,7 @@ package net.lamgc.scext.onedrive_transfer import com.microsoft.aad.msal4j.ITokenCacheAccessAspect import com.microsoft.aad.msal4j.ITokenCacheAccessContext +import mu.KotlinLogging import org.jetbrains.exposed.dao.LongEntity import org.jetbrains.exposed.dao.LongEntityClass import org.jetbrains.exposed.dao.id.EntityID @@ -28,12 +29,12 @@ class OneDriveTransferSetting(id: EntityID) : LongEntity(id) { } object TokenCaches : LongIdTable() { - val accountId = varchar("access_id", 256).uniqueIndex() + val clientId = varchar("client_id", 256).uniqueIndex() val cache = text("cache_data") } class TokenCache(id: EntityID) : LongEntity(id) { - var accountId by TokenCaches.accountId + var clientId by TokenCaches.clientId var cache by TokenCaches.cache companion object : LongEntityClass(TokenCaches) @@ -43,19 +44,22 @@ class DatabaseTokenCache(private val db: Database) : ITokenCacheAccessAspect { override fun beforeCacheAccess(context: ITokenCacheAccessContext) { transaction(db) { - TokenCache.find { TokenCaches.accountId eq context.account().homeAccountId() }.firstOrNull()?.let { + TokenCache.find { TokenCaches.clientId eq context.clientId() }.firstOrNull()?.let { context.tokenCache().deserialize(it.cache) } } } override fun afterCacheAccess(context: ITokenCacheAccessContext) { + if (!context.hasCacheChanged()) { + return + } transaction(db) { val existCache = - TokenCache.find { TokenCaches.accountId eq context.account().homeAccountId() }.firstOrNull() + TokenCache.find { TokenCaches.clientId eq context.clientId() }.firstOrNull() if (existCache == null) { TokenCache.new { - accountId = context.account().homeAccountId() + clientId = context.clientId() cache = context.tokenCache().serialize() } } else { diff --git a/src/main/kotlin/Exceptions.kt b/src/main/kotlin/Exceptions.kt new file mode 100644 index 0000000..d98fe18 --- /dev/null +++ b/src/main/kotlin/Exceptions.kt @@ -0,0 +1,3 @@ +package net.lamgc.scext.onedrive_transfer + +class OneDriveNotLoginException : Exception("OneDrive 未登录") diff --git a/src/main/kotlin/OneDriveTransferCenter.kt b/src/main/kotlin/OneDriveTransferCenter.kt index fec57e5..6f53c78 100644 --- a/src/main/kotlin/OneDriveTransferCenter.kt +++ b/src/main/kotlin/OneDriveTransferCenter.kt @@ -1,6 +1,8 @@ package net.lamgc.scext.onedrive_transfer import com.google.common.util.concurrent.AtomicDouble +import com.google.common.util.concurrent.ThreadFactoryBuilder +import com.microsoft.graph.http.GraphServiceException import com.microsoft.graph.models.DriveItem import com.microsoft.graph.models.DriveItemCreateUploadSessionParameterSet import com.microsoft.graph.models.DriveItemUploadableProperties @@ -11,44 +13,81 @@ import mu.KotlinLogging import okhttp3.Request import org.telegram.abilitybots.api.bot.BaseAbilityBot import org.telegram.telegrambots.meta.api.methods.GetFile +import org.telegram.telegrambots.meta.api.methods.send.SendMessage +import org.telegram.telegrambots.meta.api.methods.updatingmessages.EditMessageText import org.telegram.telegrambots.meta.api.objects.Document import java.util.* import java.util.concurrent.* -import kotlin.math.log object OneDriveTransferCenter { - private val queue = PriorityBlockingQueue(100, - compareBy { it.document.fileSize }.thenBy { it.createdAt.time }) + private val queue = ArrayBlockingQueue(100) + + private val executor = OneDriveTransferTaskExecutor(5, DefaultOneDriveTransferCallback, queue) fun submitUploadTask(task: OneDriveTransferTask) { - queue.put(task) + queue.offer(task) } } object DefaultOneDriveTransferCallback : OneDriveTransferCallback { + override fun onTransferStart(progress: OneDriveTransferWorkerProgress) { + val message = progress.currentTask.bot.execute( + SendMessage.builder() + .text( + """ + OneDrive 中转任务开始执行 + 文件名：${progress.currentTask.document.fileName} + """.trimIndent() + ) + .chatId(progress.currentTask.extra["chatId"].toString().toLong()) + .replyToMessageId(progress.currentTask.extra["messageId"].toString().toInt()) + .build() + ) + progress.currentTask.extra["messageId"] = message.messageId + } + override fun onProgress(progress: OneDriveTransferWorkerProgress) { - TODO("Not yet implemented") + progress.currentTask.bot.execute(EditMessageText.builder() + .chatId(progress.currentTask.extra["chatId"].toString().toLong()) + .messageId(progress.currentTask.extra["messageId"].toString().toInt()) + .text(""" + OneDrive 中转任务执行中 + 文件名：${progress.currentTask.document.fileName} + 进度：${progress.progress.get() * 100}% + """.trimIndent()) + .build()) } override fun onTransferFailure(task: OneDriveTransferTask, progress: OneDriveTransferWorkerProgress) { - task.bot.silent().send(""" - OneDrive 中转任务执行失败 - 文件名：${task.document.fileName} - """.trimIndent(), task.tgUserId) + task.bot.execute(EditMessageText.builder() + .chatId(task.extra["chatId"].toString().toLong()) + .messageId(task.extra["messageId"].toString().toInt()) + .text(""" + OneDrive 中转任务执行失败 + 文件名：${task.document.fileName} + 错误信息：${progress.exception?.message} + """.trimIndent()) + .build()) } override fun onTransferSuccess(task: OneDriveTransferTask, progress: OneDriveTransferWorkerProgress) { - task.bot.silent().send(""" - OneDrive 中转任务执行成功 - 文件名：${task.document.fileName} - """.trimIndent(), task.tgUserId) + task.bot.execute(EditMessageText.builder() + .chatId(task.extra["chatId"].toString().toLong()) + .messageId(task.extra["messageId"].toString().toInt()) + .text(""" + OneDrive 中转任务执行成功 + 文件名：${task.document.fileName} + OneDrive 文件路径：${progress.driveItem?.webUrl} + """.trimIndent()) + .build()) } } interface OneDriveTransferCallback { + fun onTransferStart(progress: OneDriveTransferWorkerProgress) fun onProgress(progress: OneDriveTransferWorkerProgress) fun onTransferFailure(task: OneDriveTransferTask, progress: OneDriveTransferWorkerProgress) fun onTransferSuccess(task: OneDriveTransferTask, progress: OneDriveTransferWorkerProgress) @@ -56,9 +95,15 @@ interface OneDriveTransferCallback { class OneDriveTransferTaskExecutor( threadNum: Int, - val callback: OneDriveTransferCallback, - val taskQueue: PriorityBlockingQueue -) : ThreadPoolExecutor(threadNum, threadNum, 0, TimeUnit.SECONDS, ArrayBlockingQueue(50)) { + private val callback: OneDriveTransferCallback, + private val taskQueue: BlockingQueue +) : ThreadPoolExecutor( + threadNum, threadNum, 0, TimeUnit.SECONDS, + ArrayBlockingQueue(50), + ThreadFactoryBuilder() + .setNameFormat("onedrive-transfer-worker-%d") + .build() +) { private val logger = KotlinLogging.logger { } @@ -66,17 +111,21 @@ class OneDriveTransferTaskExecutor( init { for (i in 0 until threadNum) { - submit(createWorker(i)) + execute(createWorker(i)) } } private fun createWorker(id: Int): Runnable = Runnable { - while (Thread.interrupted()) { + logger.info { "下载线程 $id 已启动." } + while (!Thread.interrupted()) { val task = taskQueue.take() + logger.info { "线程 $id 开始执行任务: ${task.document.fileName}" } val progress = OneDriveTransferWorkerProgress(task) threadStatusMap[id] = progress try { + callback.onTransferStart(progress) doTransferFile(task, progress) + logger.info { "OneDrive 中转任务执行成功: ${task.document.fileName}" } callback.onTransferSuccess(task, progress) } catch (e: Exception) { logger.warn(e) { "OneDrive 中转任务执行失败: ${e.message}" } @@ -84,6 +133,7 @@ class OneDriveTransferTaskExecutor( this.exception = e }) } finally { + logger.info { "线程 $id 任务执行完毕: ${task.document.fileName}" } threadStatusMap.remove(id) } } @@ -96,15 +146,14 @@ class OneDriveTransferTaskExecutor( .build() ) - val graphClient = task.service.createGraphClient(task.tgUserId) ?: throw IllegalStateException("未登录 OneDrive.") + val graphClient = task.service.createGraphClient(task.tgUserId) val drive = graphClient.drives(task.onedriveId).buildRequest().get() ?: throw IllegalStateException("无法获取 OneDrive 驱动器.") if (file.fileSize > drive.quota!!.remaining!!) { throw IllegalStateException("OneDrive 剩余空间不足.") } - // TODO: 需要完善一下文件夹位置，文件名冲突处理的问题. - graphClient.drives(task.onedriveId).root().itemWithPath(task.storagePath).buildRequest().get() - val filePath = checkAndGetPath(graphClient, task.storagePath, task.document.fileName) + val filePath = checkAndGetPath(graphClient, task.onedriveId, task.storagePath, task.document.fileName) + logger.debug { "OneDrive 中转任务: ${task.document.fileName} -> $filePath" } if (file.fileSize < 4 * 1024 * 1024) { val fileBytes = task.bot.downloadFileAsStream(file).readAllBytes() @@ -115,11 +164,9 @@ class OneDriveTransferTaskExecutor( progress.progress.set(1.0) } else { val uploadSession = graphClient.drives(task.onedriveId).root().itemWithPath(filePath) - .createUploadSession(DriveItemCreateUploadSessionParameterSet().apply { - this.item = DriveItemUploadableProperties().apply { - this.fileSize = file.fileSize - } - }) + .createUploadSession(DriveItemCreateUploadSessionParameterSet.newBuilder() + .withItem(DriveItemUploadableProperties()) + .build()) .buildRequest() .post() ?: throw IllegalStateException("无法创建 OneDrive 上传会话.") val progressCallback = IProgressCallback { current, max -> @@ -140,20 +187,109 @@ class OneDriveTransferTaskExecutor( file.fileSize, DriveItem::class.java ) - val uploadResult = largeFileUploadTask.upload(4 * 1024 * 1024, null, progressCallback) - if (progress.progress.get() == 1.0 && progress.exception != null) { - progress.driveItem = uploadResult.responseBody - } + val uploadResult = largeFileUploadTask.upload(327680 * 26, null, progressCallback) + progress.driveItem = uploadResult.responseBody } } - private fun checkAndGetPath(graphClient: GraphServiceClient, storagePath: String, fileName: String): String { - val path = if (storagePath.endsWith("/")) { - storagePath - } else { - "$storagePath/" + private fun checkAndGetPath(graphClient: GraphServiceClient, driveId: String, storagePath: String, originFileName: String): String { + val folderPath = checkAndCreateFolder(graphClient, driveId, storagePath) + val fileName = checkFileName(graphClient, driveId, folderPath, originFileName) + return "$folderPath$fileName" + } + + private fun checkAndCreateFolder(graphClient: GraphServiceClient, driveId: String, folderPath: String): String { + if (folderPath.trim() == "/") { + return "" + } + try { + val testPath = if (folderPath.startsWith('/')) { + folderPath.trimStart('/') + } else { + folderPath + } + graphClient.drives(driveId).root().itemWithPath(testPath).buildRequest().get() + logger.debug { "OneDrive 文件夹已存在：$testPath" } + return if (testPath.endsWith('/')) { + testPath + } else { + "$testPath/" + } + } catch (e: GraphServiceException) { + if (e.responseCode != 404) { + throw e + } + } + + val path = folderPath.trim('/') + val pathComponents = path.split("/") + logger.debug { "PathComponents = $pathComponents" } + var parentPath: String + var currentPath = "/" + for (component in pathComponents) { + if (component.trim().isEmpty()) { + continue + } + parentPath = currentPath + currentPath += "$component/" + logger.debug { "CurrentPath = $currentPath" } + try { + val driveItem = graphClient.drives(driveId).root().itemWithPath(currentPath).buildRequest().get() + if (driveItem!!.folder == null) { + throw IllegalStateException("OneDrive 中已存在同名文件: $currentPath") + } + } catch (e: GraphServiceException) { + if (e.responseCode == 404) { + try { + val newFolder = DriveItem() + newFolder.name = component + newFolder.folder = com.microsoft.graph.models.Folder() + graphClient.drives(driveId).root().itemWithPath(parentPath).children() + .buildRequest() + .post(newFolder) + } catch (e: GraphServiceException) { + if (e.error?.error?.code != "nameAlreadyExists") { + throw e + } + } + } else { + throw e + } + } + } + if (!currentPath.endsWith("/")) { + currentPath += "/" + } + return currentPath.trimStart('/') + } + + private fun checkFileName(graphClient: GraphServiceClient, driveId: String, folderPath: String, fileName: String): String { + try { + graphClient.drives(driveId).root().itemWithPath("$folderPath$fileName").buildRequest().get() + } catch (e: GraphServiceException) { + if (e.responseCode == 404) { + return fileName + } else { + throw e + } + } + val fileNameComponents = fileName.split(".") + val fileNameWithoutExtension = fileNameComponents.subList(0, fileNameComponents.size - 1).joinToString(".") + val fileExtension = fileNameComponents.last() + var i = 1 + while (true) { + val newFileName = "$fileNameWithoutExtension ($i).$fileExtension" + try { + graphClient.drives(driveId).root().itemWithPath("$folderPath$newFileName").buildRequest().get() + i++ + } catch (e: GraphServiceException) { + if (e.responseCode != 404) { + throw e + } + return newFileName + } + } - return "$path$fileName" } } diff --git a/src/main/kotlin/OneDriveTransferExtension.kt b/src/main/kotlin/OneDriveTransferExtension.kt index 45c7244..a5f16e8 100644 --- a/src/main/kotlin/OneDriveTransferExtension.kt +++ b/src/main/kotlin/OneDriveTransferExtension.kt @@ -2,6 +2,8 @@ package net.lamgc.scext.onedrive_transfer import com.fasterxml.jackson.databind.ObjectMapper import com.fasterxml.jackson.module.kotlin.registerKotlinModule +import com.google.common.cache.CacheBuilder +import com.microsoft.aad.msal4j.MsalInteractionRequiredException import mu.KotlinLogging import org.jetbrains.exposed.sql.Database import org.telegram.abilitybots.api.bot.BaseAbilityBot @@ -11,8 +13,18 @@ import org.telegram.abilitybots.api.objects.Locality import org.telegram.abilitybots.api.objects.Privacy import org.telegram.abilitybots.api.objects.Reply import org.telegram.abilitybots.api.util.AbilityExtension +import org.telegram.telegrambots.meta.api.methods.send.SendMessage +import org.telegram.telegrambots.meta.api.methods.updatingmessages.DeleteMessage +import org.telegram.telegrambots.meta.api.methods.updatingmessages.EditMessageText +import org.telegram.telegrambots.meta.api.objects.replykeyboard.InlineKeyboardMarkup +import org.telegram.telegrambots.meta.api.objects.replykeyboard.ReplyKeyboardMarkup +import org.telegram.telegrambots.meta.api.objects.replykeyboard.ReplyKeyboardRemove +import org.telegram.telegrambots.meta.api.objects.replykeyboard.buttons.InlineKeyboardButton +import org.telegram.telegrambots.meta.api.objects.replykeyboard.buttons.KeyboardRow import java.io.File +import java.net.MalformedURLException import java.net.URL +import java.util.concurrent.TimeUnit class OneDriveTransferExtension(val bot: BaseAbilityBot, val dataFolder: File) : AbilityExtension { @@ -20,6 +32,13 @@ class OneDriveTransferExtension(val bot: BaseAbilityBot, val dataFolder: File) : private val config: ExtensionConfig private val onedriveService: OneDriveTransferService + private val buttonTokenCache = CacheBuilder.newBuilder() + .expireAfterWrite(5, TimeUnit.MINUTES) + .build() + + private val actionCache = CacheBuilder.newBuilder() + .expireAfterWrite(5, TimeUnit.MINUTES) + .build() init { config = loadConfiguration() @@ -46,6 +65,7 @@ class OneDriveTransferExtension(val bot: BaseAbilityBot, val dataFolder: File) : .privacy(Privacy.PUBLIC) .action { ctx -> val url = onedriveService.createLoginUrl(ctx.chatId()) + actionCache.put(ctx.chatId(), "login") ctx.bot().silent().send(""" 请使用以下链接进行登录： $url @@ -54,32 +74,34 @@ class OneDriveTransferExtension(val bot: BaseAbilityBot, val dataFolder: File) : 登录成功后会显示无法访问，是正常情况，请将地址栏的链接发回给机器人。 """.trimIndent(), ctx.chatId()) } + .enableStats() .reply(Reply.of( {bot, upd -> try { val account = onedriveService.updateAccount(upd.message.chat.id, URL(upd.message.text.trim())) + actionCache.invalidate(upd.message.chatId) bot.silent().send(""" 登录成功！ Microsoft 账号：${account.userName} - 只需要向我发送文件，即可中转至 OneDrive！ + 请使用 /select_drive 选择 OneDrive 驱动器以及设置上传路径。 """.trimIndent(), upd.message.chatId) + } catch (e: MsalInteractionRequiredException) { + if (e.errorCode() == "AADSTS54005") { + bot.silent().send("登录失败，登录链接已过期，请重新登录。", upd.message.chatId) + } else { + bot.silent().send("登录失败，错误代码：${e.errorCode()}", upd.message.chatId) + } + actionCache.invalidate(upd.message.chatId) + } catch (e: MalformedURLException) { + bot.silent().send("链接格式错误，请发送正确的登录链接。", upd.message.chatId) } catch (e: Exception) { logger.error(e) { "处理 Oauth2 令牌时发生错误." } bot.silent().send("处理 Oauth2 令牌时发生错误，请稍后重试。", upd.message.chatId) + actionCache.invalidate(upd.message.chatId) } }, - { upd -> - if (!upd.hasMessage() || !upd.message.hasText().not()) { - return@of false - } - try { - URL(upd.message.text.trim()) - return@of true - } catch (e: Exception) { - bot.silent().send("链接格式错误，请重新发送。", upd.message.chatId) - return@of false - } - } + { upd -> upd.hasMessage() && upd.message.hasText() }, + { upd -> actionCache.getIfPresent(upd.message.chatId) == "login" } )) .build() @@ -96,7 +118,7 @@ class OneDriveTransferExtension(val bot: BaseAbilityBot, val dataFolder: File) : } it.bot().silent().send(""" 当前账户已登录 OneDrive. - Microsoft 账号：${account.userName} + Microsoft 账号： ${account.userName} """.trimIndent(), it.chatId()) } .build() @@ -107,16 +129,213 @@ class OneDriveTransferExtension(val bot: BaseAbilityBot, val dataFolder: File) : .locality(Locality.USER) .privacy(Privacy.PUBLIC) .action { - val currentDrive = onedriveService.getCurrentDrive(it.chatId()) + val currentDrive = try { + onedriveService.getCurrentDrive(it.chatId()) + } catch (e: OneDriveNotLoginException) { + it.bot().silent().send("当前账户没有登录 OneDrive.", it.chatId()) + return@action + } val drives = onedriveService.listDriversByUserId(it.chatId()) if (drives.isEmpty()) { it.bot().silent().send("当前账户没有 OneDrive 驱动器.", it.chatId()) return@action } + val msgContent = """ + 当前账户已登录 OneDrive. + Microsoft 账号： ${onedriveService.accountManager.getAccountByTgUserId(it.chatId())?.userName} + 当前正在使用的驱动器为：[${currentDrive?.driveType ?: "None" }] ${currentDrive?.name ?: "无"} + """.trimIndent() + + SendMessage.builder() + .disableWebPagePreview(true) + .chatId(it.chatId().toString()) + .text(msgContent) + .replyMarkup(InlineKeyboardMarkup.builder().apply { + for (drive in drives) { + keyboardRow(listOf(InlineKeyboardButton.builder() + .text("[${drive.driveType}] ${drive.name}") + .callbackData(buttonTokenCache.putToken("${drive.id}", tokenPrefix = "select_drive:")) + .build())) + } + }.build()) + .build().let { msg -> + it.bot().execute(msg) + } + } + .reply({ bot, upd -> + val driveId = buttonTokenCache.getIfPresent(upd.callbackQuery.data.substringAfter(':')) + if (driveId == null) { + bot.silent().send("无效的驱动器 ID.", upd.callbackQuery.message.chatId) + return@reply + } + val drive = try { + onedriveService.listDriversByUserId(upd.callbackQuery.from.id).firstOrNull { it.id == driveId } + } catch (e: OneDriveNotLoginException) { + bot.silent().send("当前账户没有登录 OneDrive.", upd.callbackQuery.message.chatId) + return@reply + } + if (drive == null) { + bot.silent().send("无效的驱动器 ID.", upd.callbackQuery.message.chatId) + return@reply + } + onedriveService.setDrive(upd.callbackQuery.from.id, driveId) + val editMessageText = EditMessageText.builder() + .chatId(upd.callbackQuery.message.chatId) + .messageId(upd.callbackQuery.message.messageId) + .replyMarkup( + InlineKeyboardMarkup.builder() + .clearKeyboard() + .build() + ) + .text( + """ + 已选择驱动器：[${drive.driveType}] ${drive.name} + 请使用 /set_path 设置上传路径。 + """.trimIndent() + ) + .build() + bot.execute(editMessageText) + }, { upd -> + upd.hasCallbackQuery() && upd.callbackQuery.data != null && upd.callbackQuery.data.startsWith("select_drive:") + }) + .build() + + fun setPath(): Ability = Ability.builder() + .named("set_path") + .info("设置 OneDrive 中转路径.") + .locality(Locality.USER) + .privacy(Privacy.PUBLIC) + .action { + val currentDrive = try { + onedriveService.getCurrentDrive(it.chatId()) + } catch (e: OneDriveNotLoginException) { + it.bot().silent().send("当前账户没有登录 OneDrive.", it.chatId()) + return@action + } + if (currentDrive == null) { + it.bot().silent().send("当前账户没有选择 OneDrive 驱动器.", it.chatId()) + return@action + } + val transferSetting = onedriveService.accountManager.getTransferSetting(it.chatId()) + if (transferSetting == null) { + it.bot().silent().send("当前账户没有登录 OneDrive.", it.chatId()) + return@action + } + val msgContent = """ + 当前账户已登录 OneDrive. + Microsoft 账号： ${transferSetting.userName} + 当前正在使用的驱动器为：[${currentDrive.driveType}] ${currentDrive.name} + 当前上传路径为：${transferSetting.storagePath} + ------------------------------------------------- + 请发送新的上传路径，或者发送 /cancel 取消设置。 + """.trimIndent() + actionCache.put(it.chatId(), "set_path") + SendMessage.builder() + .disableWebPagePreview(true) + .chatId(it.chatId().toString()) + .text(msgContent) + .replyMarkup(ReplyKeyboardMarkup.builder().apply { + this.oneTimeKeyboard(true) + this.selective(false) + this.inputFieldPlaceholder("OneDrive 上传路径") + this.keyboardRow(KeyboardRow().apply { + add("Telegram Files/") + }) + this.resizeKeyboard(true) + this.isPersistent(false) + }.build()) + .build().let { msg -> + it.bot().execute(msg) + } + } + .reply({ bot, upd -> + val path = upd.message.text.trim() + if (path == "/cancel") { + actionCache.invalidate(upd.message.chatId) + bot.execute(SendMessage.builder() + .chatId(upd.message.chatId.toString()) + .text("已取消设置.") + .replyMarkup(ReplyKeyboardRemove.builder() + .removeKeyboard(true) + .build()) + .build()) + return@reply + } + val transferSetting = onedriveService.accountManager.getTransferSetting(upd.message.chatId) + if (transferSetting == null) { + actionCache.invalidate(upd.message.chatId) + bot.execute(SendMessage.builder() + .chatId(upd.message.chatId.toString()) + .text("当前账户没有登录 OneDrive.") + .replyMarkup(ReplyKeyboardRemove.builder() + .removeKeyboard(true) + .build()) + .build()) + return@reply + } + onedriveService.accountManager.doSomething { + transferSetting.storagePath = path + } + actionCache.invalidate(upd.message.chatId) + bot.execute(DeleteMessage.builder() + .chatId(upd.message.chatId.toString()) + .messageId(upd.message.messageId) + .build()) + bot.execute(SendMessage.builder() + .chatId(upd.message.chatId.toString()) + .text("已设置上传路径为：$path") + .replyMarkup(ReplyKeyboardRemove.builder() + .removeKeyboard(true) + .build()) + .build()) + }, { upd -> + upd.hasMessage() && upd.message.hasText() && actionCache.getIfPresent(upd.message.chatId) == "set_path" + }) + .build() + + fun logout(): Ability = Ability.builder() + .named("logout") + .info("登出 OneDrive 账户.") + .locality(Locality.USER) + .privacy(Privacy.PUBLIC) + .action { + val transferSetting = onedriveService.accountManager.getTransferSetting(it.chatId()) + if (transferSetting == null) { + it.bot().silent().send("当前账户没有登录 OneDrive.", it.chatId()) + return@action + } + onedriveService.accountManager.doSomething { + transferSetting.delete() + } + it.bot().silent().send("已登出 OneDrive.", it.chatId()) } .build() + fun uploadDocument(): Reply = Reply.of( + { bot, upd -> + val document = upd.message.document + if (document == null) { + bot.silent().send("请发送文件.", upd.message.chatId) + return@of + } + if (document.fileSize > config.maxFileSize) { + bot.silent().send("文件大小超过限制.", upd.message.chatId) + return@of + } + onedriveService.submitUploadDocumentTask(upd.message.chatId, upd.message.messageId, document) + }, + { upd -> upd.hasMessage() && upd.message.hasDocument() }, + { upd -> + try { + val transferSetting = onedriveService.accountManager.getTransferSetting(upd.message.chatId) + transferSetting != null + } catch (e: OneDriveNotLoginException) { + false + } + } + ) + private fun AbilityBuilder.named(name: String): AbilityBuilder { return if (config.useCommandPrefix) { name("odt_$name") diff --git a/src/main/kotlin/OneDriveTransferSettingManager.kt b/src/main/kotlin/OneDriveTransferSettingManager.kt index 1f734c3..402342b 100644 --- a/src/main/kotlin/OneDriveTransferSettingManager.kt +++ b/src/main/kotlin/OneDriveTransferSettingManager.kt @@ -7,7 +7,6 @@ import com.microsoft.aad.msal4j.* import java.sql.Connection import java.net.URL import java.net.URI -import java.util.concurrent.ExecutionException class OneDriveTransferSettingManager(private val authClient: ConfidentialClientApplication, private val db: Database) { @@ -22,6 +21,12 @@ class OneDriveTransferSettingManager(private val authClient: ConfidentialClientA } } + fun doSomething(handle: () -> R): R { + return transaction(db) { + handle() + } + } + fun createAuthorizationRequest(userId: Long): URL { val parameters = AuthorizationRequestUrlParameters .builder("http://localhost:45678/", OAUTH2_SCOPE) diff --git a/src/main/kotlin/Services.kt b/src/main/kotlin/Services.kt index c55f0c5..211f925 100644 --- a/src/main/kotlin/Services.kt +++ b/src/main/kotlin/Services.kt @@ -10,8 +10,12 @@ import com.microsoft.graph.models.Drive import com.microsoft.graph.requests.GraphServiceClient 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.InetSocketAddress +import java.net.Proxy import java.net.URL import java.util.concurrent.CompletableFuture @@ -22,19 +26,23 @@ class OneDriveTransferService( private val db: Database ) { val accountManager: OneDriveTransferSettingManager - val authClient: ConfidentialClientApplication = ConfidentialClientApplication.builder( + private val authClient: ConfidentialClientApplication = ConfidentialClientApplication.builder( config.clientId, ClientCredentialFactory.createFromSecret(config.clientSecret), ) .authority(OneDriveTransferSettingManager.AUTHORITY) .setTokenCacheAccessAspect(DatabaseTokenCache(db)) + .proxy(Proxy(Proxy.Type.HTTP, InetSocketAddress.createUnresolved("127.0.0.1", 1089))) .build() init { + transaction(db) { + SchemaUtils.create(OneDriveTransferSettings, TokenCaches) + } accountManager = OneDriveTransferSettingManager(authClient, db) } - fun createGraphClient(userId: Long): GraphServiceClient? { + fun createGraphClient(userId: Long): GraphServiceClient { val cache = THREAD_CURRENT_GRAPH_CLIENT.get() if (cache?.tgUserId == userId) { return cache.client @@ -47,10 +55,7 @@ class OneDriveTransferService( GraphServiceClient.builder() .httpClient(HttpClients.createDefault(MsalAuthorizationProvider(authClient, it))) .buildClient() - } - if (serviceClient == null) { - return null - } + } ?: throw OneDriveNotLoginException() THREAD_CURRENT_GRAPH_CLIENT.set(ClientCache(userId, serviceClient)) return serviceClient } @@ -61,24 +66,52 @@ class OneDriveTransferService( fun updateAccount(userId: Long, redirectUrl: URL) = accountManager.updateAccount(userId, redirectUrl) - fun listDriversByUserId(userId: Long, page: Int = 1, size: Int = 10): List { - return createGraphClient(userId)!!.drives() + fun listDriversByUserId(userId: Long): List { + val drive = createGraphClient(userId).me().drive().buildRequest().get() + val drives = createGraphClient(userId).drives() .buildRequest() - .skip(page * size) .get()?.currentPage ?: emptyList() + return mutableListOf().apply { + if (drive != null) { + add(drive) + } + addAll(drives) + } } - fun getCurrentDrive(userId: Long): Drive { + fun setDrive(userId: Long, driveId: String) { val transferSetting = - accountManager.getTransferSetting(userId) ?: throw IllegalStateException("未登录 OneDrive.") - val graphClient = createGraphClient(userId) ?: throw IllegalStateException("未登录 OneDrive.") + 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() ?: throw IllegalStateException("无法获取当前 OneDrive 驱动器.") + .get() } - fun submitUploadDocumentTask(userId: Long, document: Document) { - + 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 { diff --git a/src/main/kotlin/Utils.kt b/src/main/kotlin/Utils.kt index 98d2def..3e766a5 100644 --- a/src/main/kotlin/Utils.kt +++ b/src/main/kotlin/Utils.kt @@ -1,5 +1,6 @@ package net.lamgc.scext.onedrive_transfer +import com.google.common.cache.Cache import java.net.URL import java.net.URLDecoder @@ -11,3 +12,18 @@ fun URL.getQueryMap(): Map { } return queryMap } + +// 生成大小写字母加数字的随机字符串 +fun randomString(length: Int): String { + val charPool = ('a'..'z') + ('A'..'Z') + ('0'..'9') + return (1..length) + .map { kotlin.random.Random.nextInt(0, charPool.size) } + .map(charPool::get) + .joinToString("") +} + +fun Cache.putToken(value: String, tokenLength: Int = 16, tokenPrefix: String = ""): String { + val token = randomString(tokenLength) + put(token, value) + return "$tokenPrefix$token" +}