From d9d477f954cae46cd4ab86025870d73ad0995e12 Mon Sep 17 00:00:00 2001
From: lifegpc <g1710431395@gmail.com>
Date: Sun, 25 Sep 2022 05:26:38 +0000
Subject: [PATCH] Update

---
 src/db/sqlite/db.rs         | 26 ++++++++++++++++++++++++++
 src/db/traits.rs            | 11 +++++++++++
 src/db/user.rs              | 13 +++++++++++++
 src/server/auth/user.rs     | 13 ++++++++++++-
 src/server/unittest/auth.rs | 14 ++++++++++----
 5 files changed, 72 insertions(+), 5 deletions(-)

diff --git a/src/db/sqlite/db.rs b/src/db/sqlite/db.rs
index 130154c..311723a 100644
--- a/src/db/sqlite/db.rs
+++ b/src/db/sqlite/db.rs
@@ -84,6 +84,21 @@ pub struct PixivDownloaderSqlite {
 }
 
 impl PixivDownloaderSqlite {
+    async fn _add_root_user(
+        &self,
+        name: &str,
+        username: &str,
+        password: &[u8],
+    ) -> Result<(), PixivDownloaderDbError> {
+        let mut db = self.db.lock().await;
+        let tx = db.transaction()?;
+        tx.execute(
+            "INSERT OR REPLACE INTO users VALUES (0, ?, ?, ?, true);",
+            (name, username, password),
+        )?;
+        tx.commit()?;
+        Ok(())
+    }
     /// Check if the database needed create all tables.
     async fn _check_database(&self) -> Result<bool, SqliteError> {
         let tables = self._get_exists_table().await?;
@@ -241,6 +256,17 @@ impl PixivDownloaderDb for PixivDownloaderSqlite {
         }
     }
 
+    #[cfg(feature = "server")]
+    async fn add_root_user(
+        &self,
+        name: &str,
+        username: &str,
+        password: &[u8],
+    ) -> Result<User, PixivDownloaderDbError> {
+        self._add_root_user(name, username, password).await?;
+        Ok(self.get_user(0).await?.expect("Root user not found:"))
+    }
+
     #[cfg(feature = "server")]
     async fn get_user(&self, id: u64) -> Result<Option<User>, PixivDownloaderDbError> {
         Ok(self._get_user(id).await?)
diff --git a/src/db/traits.rs b/src/db/traits.rs
index 598813c..06cf018 100644
--- a/src/db/traits.rs
+++ b/src/db/traits.rs
@@ -13,6 +13,17 @@ pub trait PixivDownloaderDb {
     where
         Self: Sized + Send + Sync;
     #[cfg(feature = "server")]
+    /// Add root user to database.
+    /// * `name` - User name
+    /// * `username` - Unique user name
+    /// * `password` - Hashed password
+    async fn add_root_user(
+        &self,
+        name: &str,
+        username: &str,
+        password: &[u8],
+    ) -> Result<User, PixivDownloaderDbError>;
+    #[cfg(feature = "server")]
     /// Get a user by ID
     /// * `id`: The user's ID
     async fn get_user(&self, id: u64) -> Result<Option<User>, PixivDownloaderDbError>;
diff --git a/src/db/user.rs b/src/db/user.rs
index fe99be9..8b03898 100644
--- a/src/db/user.rs
+++ b/src/db/user.rs
@@ -1,5 +1,7 @@
 use bytes::BytesMut;
 
+use crate::ext::json::ToJson2;
+
 /// A user in the database
 pub struct User {
     /// The user ID
@@ -13,3 +15,14 @@ pub struct User {
     /// Whether the user is an admin
     pub is_admin: bool,
 }
+
+impl ToJson2 for User {
+    fn to_json2(&self) -> json::JsonValue {
+        json::object! {
+            "id": self.id,
+            "name": self.name.as_str(),
+            "username": self.username.as_str(),
+            "is_admin": self.is_admin,
+        }
+    }
+}
diff --git a/src/server/auth/user.rs b/src/server/auth/user.rs
index ab4e04c..f174fa0 100644
--- a/src/server/auth/user.rs
+++ b/src/server/auth/user.rs
@@ -58,10 +58,21 @@ impl AuthUserContext {
                             let password = key
                                 .decrypt(&password)
                                 .try_err3(7, gettext("Failed to decrypt password with RSA:"))?;
+                            let hashed_password = openssl::sha::sha512(&password);
+                            if root_user.is_none() {
+                                let user = self
+                                    .ctx
+                                    .db
+                                    .add_root_user(name, username, &hashed_password)
+                                    .await
+                                    .try_err3(8, gettext("Failed to add user to database:"))?;
+                                Ok(user.to_json2())
+                            } else {
+                                Ok(json::object! {})
+                            }
                         }
                         None => return Err((5, gettext("No RSA key found.")).into()),
                     }
-                    Ok(json::object! {})
                 }
             },
             None => {
diff --git a/src/server/unittest/auth.rs b/src/server/unittest/auth.rs
index e5a3d9e..2566400 100644
--- a/src/server/unittest/auth.rs
+++ b/src/server/unittest/auth.rs
@@ -4,10 +4,7 @@ use crate::ext::json::FromJson;
 use crate::server::result::JSONResult;
 use bytes::BytesMut;
 use hyper::{Body, Request};
-use openssl::{
-    pkey::Public,
-    rsa::{Padding, Rsa},
-};
+use openssl::rsa::{Padding, Rsa};
 
 /// Test authentification methods
 /// Returns token
@@ -40,5 +37,14 @@ pub async fn test(ctx: &UnitTestContext) -> Result<BytesMut, PixivDownloaderErro
         .await?
         .unwrap();
     let result = JSONResult::from_json(re)?.expect("Failed to add user:");
+    assert_eq!(
+        result,
+        json::object! {
+            "id": 0,
+            "name": "test",
+            "username": "test",
+            "is_admin": true,
+        }
+    );
     Ok(BytesMut::new())
 }