LamGC/scext-onedrive-transfer
1
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 看来基本完工了！

Browse Source
This commit is contained in:
LamGC 2024-01-10 00:17:23 +08:00
parent 4e07309066
commit 54a17521f0
Signed by: LamGC
GPG Key ID: 6C5AE2A913941E1D
8 changed files with 489 additions and 74 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
build.gradle.kts
View File

@ -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")

14
src/main/kotlin/Databases.kt
View File

@ -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<Long>) : 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<Long>) : LongEntity(id) {
var accountId by TokenCaches.accountId
var clientId by TokenCaches.clientId
var cache by TokenCaches.cache
companion object : LongEntityClass<TokenCache>(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 {

3
src/main/kotlin/Exceptions.kt Normal file
View File

@ -0,0 +1,3 @@
package net.lamgc.scext.onedrive_transfer
class OneDriveNotLoginException : Exception("OneDrive 未登录")

210
src/main/kotlin/OneDriveTransferCenter.kt
View File

@ -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<OneDriveTransferTask> { it.document.fileSize }.thenBy { it.createdAt.time })
private val queue = ArrayBlockingQueue<OneDriveTransferTask>(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<OneDriveTransferTask>
) : ThreadPoolExecutor(threadNum, threadNum, 0, TimeUnit.SECONDS, ArrayBlockingQueue(50)) {
private val callback: OneDriveTransferCallback,
private val taskQueue: BlockingQueue<OneDriveTransferTask>
) : 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<Request>, storagePath: String, fileName: String): String {
val path = if (storagePath.endsWith("/")) {
storagePath
} else {
"$storagePath/"
private fun checkAndGetPath(graphClient: GraphServiceClient<Request>, 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<Request>, 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<Request>, 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"
}
}

249
src/main/kotlin/OneDriveTransferExtension.kt
View File

@ -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<String, String>()
private val actionCache = CacheBuilder.newBuilder()
.expireAfterWrite(5, TimeUnit.MINUTES)
.build<Long, String>()
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")

7
src/main/kotlin/OneDriveTransferSettingManager.kt
View File

@ -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 <R> doSomething(handle: () -> R): R {
return transaction(db) {
handle()
}
}
fun createAuthorizationRequest(userId: Long): URL {
val parameters = AuthorizationRequestUrlParameters
.builder("http://localhost:45678/", OAUTH2_SCOPE)

63
src/main/kotlin/Services.kt
View File

@ -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<Request>? {
fun createGraphClient(userId: Long): GraphServiceClient<Request> {
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<Drive> {
return createGraphClient(userId)!!.drives()
fun listDriversByUserId(userId: Long): List<Drive> {
val drive = createGraphClient(userId).me().drive().buildRequest().get()
val drives = createGraphClient(userId).drives()
.buildRequest()
.skip(page * size)
.get()?.currentPage ?: emptyList()
return mutableListOf<Drive>().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 {

16
src/main/kotlin/Utils.kt
View File

@ -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<String, String> {
}
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<String, String>.putToken(value: String, tokenLength: Int = 16, tokenPrefix: String = ""): String {
val token = randomString(tokenLength)
put(token, value)
return "$tokenPrefix$token"
}