From 7bdb15208d180d548bdf6d867260b5ab375e2ffb Mon Sep 17 00:00:00 2001
From: Christopher Ng <chrng8@gmail.com>
Date: Tue, 1 Mar 2022 02:44:30 +0000
Subject: [PATCH] Avatar export and import

Signed-off-by: Christopher Ng <chrng8@gmail.com>
---
 lib/private/Accounts/AccountMigrator.php | 52 +++++++++++++++++++++++-
 1 file changed, 50 insertions(+), 2 deletions(-)

diff --git a/lib/private/Accounts/AccountMigrator.php b/lib/private/Accounts/AccountMigrator.php
index 3f684ce860a..b4ad375a99a 100644
--- a/lib/private/Accounts/AccountMigrator.php
+++ b/lib/private/Accounts/AccountMigrator.php
@@ -27,28 +27,45 @@ declare(strict_types=1);
 namespace OC\Accounts;
 
 use InvalidArgumentException;
+use OC\NotSquareException;
 use OCP\Accounts\IAccountManager;
+use OCP\IAvatarManager;
 use OCP\IUser;
 use OCP\UserMigration\IExportDestination;
 use OCP\UserMigration\IImportSource;
 use OCP\UserMigration\IMigrator;
 use OCP\UserMigration\TMigratorBasicVersionHandling;
 use Symfony\Component\Console\Output\OutputInterface;
+use Throwable;
 
 class AccountMigrator implements IMigrator {
-
 	use TMigratorBasicVersionHandling;
 
 	use TAccountsHelper;
 
 	private IAccountManager $accountManager;
 
+	private IAvatarManager $avatarManager;
+
 	private const EXPORT_FILE = 'account.json';
 
 	public function __construct(
-		IAccountManager $accountManager
+		IAccountManager $accountManager,
+		IAvatarManager $avatarManager
 	) {
 		$this->accountManager = $accountManager;
+		$this->avatarManager = $avatarManager;
+	}
+
+	private function getExtension(string $mimeType): string {
+		switch ($mimeType) {
+			case 'image/jpeg':
+				return 'jpg';
+			case 'image/png':
+				return 'png';
+			default:
+				throw new AccountMigratorException("Invalid avatar mimetype: \"$mimeType\"");
+		}
 	}
 
 	/**
@@ -60,6 +77,16 @@ class AccountMigrator implements IMigrator {
 		if ($exportDestination->addFileContents(AccountMigrator::EXPORT_FILE, json_encode($this->accountManager->getAccount($user))) === false) {
 			throw new AccountMigratorException('Could not export account information');
 		}
+
+		$avatar = $this->avatarManager->getAvatar($user->getUID());
+		if ($avatar->isCustomAvatar()) {
+			$avatarData = $avatar->get(-1)->data();
+			$ext = $this->getExtension($avatar->get(-1)->dataMimeType());
+			$output->writeln('Exporting avatar to avatar.' . $ext . '…');
+			if ($exportDestination->addFileContents("avatar.$ext", $avatarData) === false) {
+				throw new AccountMigratorException('Could not export avatar');
+			}
+		}
 	}
 
 	/**
@@ -99,5 +126,26 @@ class AccountMigrator implements IMigrator {
 		} catch (InvalidArgumentException $e) {
 			throw new AccountMigratorException('Failed to import account information');
 		}
+
+		foreach ($importSource->getFolderListing('') as $filename) {
+			if (str_starts_with($filename, 'avatar.')) {
+				$avatarFilename = $filename;
+			}
+		}
+
+		if (isset($avatarFilename)) {
+			$output->writeln('Importing avatar from ' . $avatarFilename . '…');
+			$avatar = $importSource->getFileContents($avatarFilename);
+			$image = new \OC_Image();
+			$image->loadFromData($avatar);
+			try {
+				$avatar = $this->avatarManager->getAvatar($user->getUID());
+				$avatar->set($image);
+			} catch (NotSquareException $e) {
+				throw new AccountMigratorException('Avatar image must be square');
+			} catch (Throwable $e) {
+				throw new AccountMigratorException('Failed to import avatar', 0, $e);
+			}
+		}
 	}
 }