5 files changed, 156 insertions, 25 deletions

diff --git a/shared/server-commands/server/server.ts b/shared/server-commands/server/server.ts
index a8f8c1d84..7096faf21 100644
--- a/shared/server-commands/server/server.ts
+++ b/shared/server-commands/server/server.ts
@@ -13,7 +13,15 @@ import { AbusesCommand } from '../moderation'
 import { OverviewsCommand } from '../overviews'
 import { SearchCommand } from '../search'
 import { SocketIOCommand } from '../socket'
-import { AccountsCommand, BlocklistCommand, LoginCommand, NotificationsCommand, SubscriptionsCommand, UsersCommand } from '../users'
+import {
+  AccountsCommand,
+  BlocklistCommand,
+  LoginCommand,
+  NotificationsCommand,
+  SubscriptionsCommand,
+  TwoFactorCommand,
+  UsersCommand
+} from '../users'
 import {
   BlacklistCommand,
   CaptionsCommand,
@@ -136,6 +144,7 @@ export class PeerTubeServer {
   videos?: VideosCommand
   videoStats?: VideoStatsCommand
   views?: ViewsCommand
+  twoFactor?: TwoFactorCommand
 
   constructor (options: { serverNumber: number } | { url: string }) {
     if ((options as any).url) {
@@ -417,5 +426,6 @@ export class PeerTubeServer {
     this.videoStudio = new VideoStudioCommand(this)
     this.videoStats = new VideoStatsCommand(this)
     this.views = new ViewsCommand(this)
+    this.twoFactor = new TwoFactorCommand(this)
   }
 }
diff --git a/shared/server-commands/users/index.ts b/shared/server-commands/users/index.ts
index f6f93b4d2..1afc02dc1 100644
--- a/shared/server-commands/users/index.ts
+++ b/shared/server-commands/users/index.ts
@@ -5,4 +5,5 @@ export * from './login'
 export * from './login-command'
 export * from './notifications-command'
 export * from './subscriptions-command'
+export * from './two-factor-command'
 export * from './users-command'
diff --git a/shared/server-commands/users/login-command.ts b/shared/server-commands/users/login-command.ts
index 54070e426..f2fc6d1c5 100644
--- a/shared/server-commands/users/login-command.ts
+++ b/shared/server-commands/users/login-command.ts
@@ -2,34 +2,27 @@ import { HttpStatusCode, PeerTubeProblemDocument } from '@shared/models'
 import { unwrapBody } from '../requests'
 import { AbstractCommand, OverrideCommandOptions } from '../shared'
 
+type LoginOptions = OverrideCommandOptions & {
+  client?: { id?: string, secret?: string }
+  user?: { username: string, password?: string }
+  otpToken?: string
+}
+
 export class LoginCommand extends AbstractCommand {
 
-  login (options: OverrideCommandOptions & {
-    client?: { id?: string, secret?: string }
-    user?: { username: string, password?: string }
-  } = {}) {
-    const { client = this.server.store.client, user = this.server.store.user } = options
-    const path = '/api/v1/users/token'
+  async login (options: LoginOptions = {}) {
+    const res = await this._login(options)
 
-    const body = {
-      client_id: client.id,
-      client_secret: client.secret,
-      username: user.username,
-      password: user.password ?? 'password',
-      response_type: 'code',
-      grant_type: 'password',
-      scope: 'upload'
-    }
+    return this.unwrapLoginBody(res.body)
+  }
 
-    return unwrapBody<{ access_token: string, refresh_token: string } & PeerTubeProblemDocument>(this.postBodyRequest({
-      ...options,
+  async loginAndGetResponse (options: LoginOptions = {}) {
+    const res = await this._login(options)
 
-      path,
-      requestType: 'form',
-      fields: body,
-      implicitToken: false,
-      defaultExpectedStatus: HttpStatusCode.OK_200
-    }))
+    return {
+      res,
+      body: this.unwrapLoginBody(res.body)
+    }
   }
 
   getAccessToken (arg1?: { username: string, password?: string }): Promise<string>
@@ -129,4 +122,38 @@ export class LoginCommand extends AbstractCommand {
       defaultExpectedStatus: HttpStatusCode.OK_200
     })
   }
+
+  private _login (options: LoginOptions) {
+    const { client = this.server.store.client, user = this.server.store.user, otpToken } = options
+    const path = '/api/v1/users/token'
+
+    const body = {
+      client_id: client.id,
+      client_secret: client.secret,
+      username: user.username,
+      password: user.password ?? 'password',
+      response_type: 'code',
+      grant_type: 'password',
+      scope: 'upload'
+    }
+
+    const headers = otpToken
+      ? { 'x-peertube-otp': otpToken }
+      : {}
+
+    return this.postBodyRequest({
+      ...options,
+
+      path,
+      headers,
+      requestType: 'form',
+      fields: body,
+      implicitToken: false,
+      defaultExpectedStatus: HttpStatusCode.OK_200
+    })
+  }
+
+  private unwrapLoginBody (body: any) {
+    return body as { access_token: string, refresh_token: string } & PeerTubeProblemDocument
+  }
 }
diff --git a/shared/server-commands/users/two-factor-command.ts b/shared/server-commands/users/two-factor-command.ts
new file mode 100644
index 000000000..5542acfda
--- /dev/null
+++ b/shared/server-commands/users/two-factor-command.ts
@@ -0,0 +1,92 @@
+import { TOTP } from 'otpauth'
+import { HttpStatusCode, TwoFactorEnableResult } from '@shared/models'
+import { unwrapBody } from '../requests'
+import { AbstractCommand, OverrideCommandOptions } from '../shared'
+
+export class TwoFactorCommand extends AbstractCommand {
+
+  static buildOTP (options: {
+    secret: string
+  }) {
+    const { secret } = options
+
+    return new TOTP({
+      issuer: 'PeerTube',
+      algorithm: 'SHA1',
+      digits: 6,
+      period: 30,
+      secret
+    })
+  }
+
+  request (options: OverrideCommandOptions & {
+    userId: number
+    currentPassword?: string
+  }) {
+    const { currentPassword, userId } = options
+
+    const path = '/api/v1/users/' + userId + '/two-factor/request'
+
+    return unwrapBody<TwoFactorEnableResult>(this.postBodyRequest({
+      ...options,
+
+      path,
+      fields: { currentPassword },
+      implicitToken: true,
+      defaultExpectedStatus: HttpStatusCode.OK_200
+    }))
+  }
+
+  confirmRequest (options: OverrideCommandOptions & {
+    userId: number
+    requestToken: string
+    otpToken: string
+  }) {
+    const { userId, requestToken, otpToken } = options
+
+    const path = '/api/v1/users/' + userId + '/two-factor/confirm-request'
+
+    return this.postBodyRequest({
+      ...options,
+
+      path,
+      fields: { requestToken, otpToken },
+      implicitToken: true,
+      defaultExpectedStatus: HttpStatusCode.NO_CONTENT_204
+    })
+  }
+
+  disable (options: OverrideCommandOptions & {
+    userId: number
+    currentPassword?: string
+  }) {
+    const { userId, currentPassword } = options
+    const path = '/api/v1/users/' + userId + '/two-factor/disable'
+
+    return this.postBodyRequest({
+      ...options,
+
+      path,
+      fields: { currentPassword },
+      implicitToken: true,
+      defaultExpectedStatus: HttpStatusCode.NO_CONTENT_204
+    })
+  }
+
+  async requestAndConfirm (options: OverrideCommandOptions & {
+    userId: number
+    currentPassword?: string
+  }) {
+    const { userId, currentPassword } = options
+
+    const { otpRequest } = await this.request({ userId, currentPassword })
+
+    await this.confirmRequest({
+      userId,
+      requestToken: otpRequest.requestToken,
+      otpToken: TwoFactorCommand.buildOTP({ secret: otpRequest.secret }).generate()
+    })
+
+    return otpRequest
+  }
+}
diff --git a/shared/server-commands/users/users-command.ts b/shared/server-commands/users/users-command.ts
index e7d021059..811b9685b 100644
--- a/shared/server-commands/users/users-command.ts
+++ b/shared/server-commands/users/users-command.ts
@@ -202,7 +202,8 @@ export class UsersCommand extends AbstractCommand {
       token,
       userId: user.id,
       userChannelId: me.videoChannels[0].id,
-      userChannelName: me.videoChannels[0].name
+      userChannelName: me.videoChannels[0].name,
+      password
     }
   }