4 files changed, 128 insertions, 24 deletions

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..6c9d270ae
--- /dev/null
+++ b/shared/server-commands/users/two-factor-command.ts
@@ -0,0 +1,75 @@
+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
+    })
+  }
+}
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
     }
   }