Redo migration table

src/router/userdata/user/migration.rs

@@ -5,26 +5,54 @@ use rand::Rng;
 use sha2::{Digest, Sha256};
 use base64::{Engine as _, engine::general_purpose};
 
-pub fn get_acc_transfer(uid: i64, token: &str, password: &str) -> JsonValue {
+fn generate_token() -> String {
+    let charset = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
+    let mut rng = rand::rng();
+    let random_string: String = (0..16)
+        .map(|_| {
+            let idx = rng.random_range(0..charset.len());
+            charset.chars().nth(idx).unwrap()
+        })
+        .collect();
+    random_string
+}
+
+pub fn get_acc_transfer(token: &str, password: &str) -> JsonValue {
     let database = userdata::get_userdata_database();
     let data = database.lock_and_select("SELECT password FROM migration WHERE token=?1", params!(token));
     if data.is_err() {
         return object!{success: false};
     }
     if verify_password(password, &data.unwrap()) {
+        let uid: i64 = database.lock_and_select_type("SELECT user_id FROM migration WHERE token=?1", params!(token)).unwrap();
         let login_token = userdata::get_login_token(uid);
         if login_token == String::new() {
             return object!{success: false};
         }
-        return object!{success: true, login_token: login_token};
+        return object!{success: true, login_token: login_token, user_id: uid};
     }
     object!{success: false}
 }
 
-pub fn save_acc_transfer(token: &str, password: &str) {
+pub fn save_acc_transfer(uid: i64, password: &str) -> String {
     let database = userdata::get_userdata_database();
-    database.lock_and_exec("DELETE FROM migration WHERE token=?1", params!(token));
-    database.lock_and_exec("INSERT INTO migration (token, password) VALUES (?1, ?2)", params!(token, hash_password(password)));
+    let token = if let Ok(value) = database.lock_and_select("SELECT token FROM migration WHERE user_id=?1", params!(uid)) {
+        value
+    } else {
+        generate_token()
+    };
+    database.lock_and_exec("DELETE FROM migration WHERE user_id=?1", params!(uid));
+    database.lock_and_exec("INSERT INTO migration (user_id, token, password) VALUES (?1, ?2, ?3)", params!(uid, &token, hash_password(password)));
+    token
+}
+
+pub fn get_acc_token(uid: i64) -> String {
+    let database = userdata::get_userdata_database();
+    if let Ok(value) = database.lock_and_select("SELECT token FROM migration WHERE user_id=?1", params!(uid)) {
+        value
+    } else {
+        save_acc_transfer(uid, "")
+    }
 }
 
 fn generate_salt() -> Vec<u8> {
@@ -35,6 +63,7 @@ fn generate_salt() -> Vec<u8> {
 }
 
 fn hash_password(password: &str) -> String {
+    if password.is_empty() { return String::new(); };
     let salt = &generate_salt();
     let mut hasher = Sha256::new();
     hasher.update(password.as_bytes());
@@ -67,3 +96,62 @@ pub fn verify_password(password: &str, salted_hash: &str) -> bool {
 
     input_hash.as_slice() == hashed_password
 }
+
+pub fn setup_sql(conn: &rusqlite::Connection) -> Result<(), rusqlite::Error> {
+    conn.execute("
+    CREATE TABLE IF NOT EXISTS migration (
+        user_id           BIGINT NOT NULL,
+        token             TEXT NOT NULL,
+        password          TEXT NOT NULL,
+        PRIMARY KEY (user_id, token)
+    );
+    ", [])?;
+    let is_updated = conn.prepare("SELECT user_id FROM migration LIMIT 1;").is_ok();
+    if is_updated { return Ok(()); }
+    println!("Upgrading migration table");
+    conn.execute("DROP TABLE IF EXISTS migration_new;", [])?;
+    conn.execute("
+    CREATE TABLE migration_new (
+        user_id           BIGINT NOT NULL,
+        token             TEXT NOT NULL,
+        password          TEXT NOT NULL,
+        PRIMARY KEY (user_id, token)
+    );
+    ", [])?;
+    
+    let mut stmt = conn.prepare("SELECT token, password FROM migration")?;
+
+    let rows = stmt.query_map([], |row| {
+        Ok((
+            row.get::<_, String>(0)?,
+            row.get::<_, String>(1)?,
+        ))
+    })?;
+
+    let mut insert_new_row = conn.prepare("INSERT INTO migration_new (user_id, token, password) VALUES (?1, ?2, ?3)")?;
+    for row in rows {
+        let (token, password) = row?;
+
+        let user_id = code_to_uid(&token);
+
+        insert_new_row.execute(params![user_id, token, password])?;
+    }
+    conn.execute("DROP TABLE migration;", params!())?;
+    conn.execute("ALTER TABLE migration_new RENAME TO migration;", params!())?;
+    Ok(())
+}
+
+fn code_to_uid(code: &str) -> String {
+    code
+        .replace('7', "")
+        .replace('A', "1")
+        .replace('G', "2")
+        .replace('W', "3")
+        .replace('Q', "4")
+        .replace('Y', "5")
+        .replace('6', "6")
+        .replace('I', "7")
+        .replace('P', "8")
+        .replace('U', "9")
+        .replace('M', "0")
+}