From b0ee5f65703d3f89c7511a8700cfd212b95b36fc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Garc=C3=ADa?=
 <dani-garcia@users.noreply.github.com>
Date: Fri, 1 Jun 2018 15:08:03 +0200
Subject: [PATCH] Improved two factor auth

---
 .../down.sql                                  |   1 +
 .../up.sql                                    |   3 +
 src/api/core/mod.rs                           |  43 +++-
 src/api/core/two_factor.rs                    |   2 +-
 src/api/identity.rs                           | 230 ++++++++++++------
 src/db/models/device.rs                       |  15 ++
 src/db/models/user.rs                         |  35 +--
 src/db/schema.rs                              |  21 +-
 8 files changed, 247 insertions(+), 103 deletions(-)
 create mode 100644 migrations/2018-06-01-112529_update_devices_twofactor_remember/down.sql
 create mode 100644 migrations/2018-06-01-112529_update_devices_twofactor_remember/up.sql

diff --git a/migrations/2018-06-01-112529_update_devices_twofactor_remember/down.sql b/migrations/2018-06-01-112529_update_devices_twofactor_remember/down.sql
new file mode 100644
index 00000000..291a97c5
--- /dev/null
+++ b/migrations/2018-06-01-112529_update_devices_twofactor_remember/down.sql
@@ -0,0 +1 @@
+-- This file should undo anything in `up.sql`
\ No newline at end of file
diff --git a/migrations/2018-06-01-112529_update_devices_twofactor_remember/up.sql b/migrations/2018-06-01-112529_update_devices_twofactor_remember/up.sql
new file mode 100644
index 00000000..aaad8eab
--- /dev/null
+++ b/migrations/2018-06-01-112529_update_devices_twofactor_remember/up.sql
@@ -0,0 +1,3 @@
+ALTER TABLE devices
+    ADD COLUMN
+    twofactor_remember TEXT;
\ No newline at end of file
diff --git a/src/api/core/mod.rs b/src/api/core/mod.rs
index 34cf9472..99a8a523 100644
--- a/src/api/core/mod.rs
+++ b/src/api/core/mod.rs
@@ -96,22 +96,49 @@ pub fn routes() -> Vec<Route> {
 
 use rocket::Route;
 
-use rocket_contrib::Json;
+use rocket_contrib::{Json, Value};
 
 use db::DbConn;
+use db::models::*;
 
 use api::{JsonResult, EmptyResult, JsonUpcase};
 use auth::Headers;
 
-#[put("/devices/identifier/<uuid>/clear-token")]
-fn clear_device_token(uuid: String, _conn: DbConn) -> JsonResult {
-    println!("{}", uuid);
-    err!("Not implemented")
+#[put("/devices/identifier/<uuid>/clear-token", data = "<data>")]
+fn clear_device_token(uuid: String, data: Json<Value>, headers: Headers, conn: DbConn) -> EmptyResult {
+    println!("UUID: {:#?}", uuid);
+    println!("DATA: {:#?}", data);
+    
+    let device = match Device::find_by_uuid(&uuid, &conn) {
+        Some(device) => device,
+        None => err!("Device not found")
+    };
+
+    if device.user_uuid != headers.user.uuid {
+        err!("Device not owned by user")
+    }
+
+    device.delete(&conn);
+
+    Ok(())
 }
 
-#[put("/devices/identifier/<uuid>/token")]
-fn put_device_token(uuid: String, _conn: DbConn) -> JsonResult {
-    println!("{}", uuid);
+#[put("/devices/identifier/<uuid>/token", data = "<data>")]
+fn put_device_token(uuid: String, data: Json<Value>, headers: Headers, conn: DbConn) -> JsonResult {
+    println!("UUID: {:#?}", uuid);
+    println!("DATA: {:#?}", data);
+
+    let device = match Device::find_by_uuid(&uuid, &conn) {
+        Some(device) => device,
+        None => err!("Device not found")
+    };
+
+    if device.user_uuid != headers.user.uuid {
+        err!("Device not owned by user")
+    }
+
+    // TODO: What does this do?
+
     err!("Not implemented")
 }
 
diff --git a/src/api/core/two_factor.rs b/src/api/core/two_factor.rs
index 7a5856f7..c226519b 100644
--- a/src/api/core/two_factor.rs
+++ b/src/api/core/two_factor.rs
@@ -135,7 +135,7 @@ fn activate_authenticator(data: JsonUpcase<EnableTwoFactorData>, headers: Header
     user.totp_secret = Some(key.to_uppercase());
 
     // Validate the token provided with the key
-    if !user.check_totp_code(Some(token)) {
+    if !user.check_totp_code(token) {
         err!("Invalid totp code")
     }
 
diff --git a/src/api/identity.rs b/src/api/identity.rs
index 805c334f..df035e48 100644
--- a/src/api/identity.rs
+++ b/src/api/identity.rs
@@ -3,7 +3,7 @@ use std::collections::HashMap;
 use rocket::{Route, Outcome};
 use rocket::request::{self, Request, FromRequest, Form, FormItems, FromForm};
 
-use rocket_contrib::Json;
+use rocket_contrib::{Json, Value};
 
 use db::DbConn;
 use db::models::*;
@@ -19,98 +19,192 @@ pub fn routes() -> Vec<Route> {
 #[post("/connect/token", data = "<connect_data>")]
 fn login(connect_data: Form<ConnectData>, device_type: DeviceType, conn: DbConn) -> JsonResult {
     let data = connect_data.get();
+    println!("{:#?}", data);
 
-    let mut device = match data.grant_type {
-        GrantType::RefreshToken => {
-            // Extract token
-            let token = data.get("refresh_token").unwrap();
+    match data.grant_type {
+        GrantType::RefreshToken =>_refresh_login(data, device_type, conn),
+        GrantType::Password => _password_login(data, device_type, conn)
+    }
+}
 
-            // Get device by refresh token
-            match Device::find_by_refresh_token(token, &conn) {
-                Some(device) => device,
-                None => err!("Invalid refresh token")
-            }
-        }
-        GrantType::Password => {
-            // Validate scope
-            let scope = data.get("scope").unwrap();
-            if scope != "api offline_access" {
-                err!("Scope not supported")
-            }
+fn _refresh_login(data: &ConnectData, _device_type: DeviceType, conn: DbConn) -> JsonResult {
+    // Extract token
+    let token = data.get("refresh_token").unwrap();
 
-            // Get the user
-            let username = data.get("username").unwrap();
-            let user = match User::find_by_mail(username, &conn) {
-                Some(user) => user,
-                None => err!("Username or password is incorrect. Try again.")
-            };
+    // Get device by refresh token
+    let mut device = match Device::find_by_refresh_token(token, &conn) {
+        Some(device) => device,
+        None => err!("Invalid refresh token")
+    };
 
-            // Check password
-            let password = data.get("password").unwrap();
-            if !user.check_valid_password(password) {
-                err!("Username or password is incorrect. Try again.")
-            }
+    // COMMON
+    let user = User::find_by_uuid(&device.user_uuid, &conn).unwrap();
+    let orgs = UserOrganization::find_by_user(&user.uuid, &conn);
 
-            // Check if totp code is required and the value is correct
-            let totp_code = util::parse_option_string(data.get("twoFactorToken"));
-
-            if !user.check_totp_code(totp_code) {
-                // Return error 400
-                err_json!(json!({
-                    "error" : "invalid_grant",
-                    "error_description" : "Two factor required.",
-                    "TwoFactorProviders" : [ 0 ],
-                    "TwoFactorProviders2" : { "0" : null }
-                }))
-            }
+    let (access_token, expires_in) = device.refresh_tokens(&user, orgs);
+    device.save(&conn);
 
-            // Let's only use the header and ignore the 'devicetype' parameter
-            let device_type_num = device_type.0;
+    Ok(Json(json!({
+        "access_token": access_token,
+        "expires_in": expires_in,
+        "token_type": "Bearer",
+        "refresh_token": device.refresh_token,
+        "Key": user.key,
+        "PrivateKey": user.private_key,
+    })))
+}
 
-            let (device_id, device_name) = match data.is_device {
-                false => { (format!("web-{}", user.uuid), String::from("web")) }
-                true => {
-                    (
-                        data.get("deviceidentifier").unwrap().clone(),
-                        data.get("devicename").unwrap().clone(),
-                    )
-                }
-            };
+fn _password_login(data: &ConnectData, device_type: DeviceType, conn: DbConn) -> JsonResult {
+    // Validate scope
+    let scope = data.get("scope").unwrap();
+    if scope != "api offline_access" {
+        err!("Scope not supported")
+    }
 
-            // Find device or create new
-            match Device::find_by_uuid(&device_id, &conn) {
-                Some(device) => {
-                    // Check if valid device
-                    if device.user_uuid != user.uuid {
-                        device.delete(&conn);
-                        err!("Device is not owned by user")
-                    }
+    // Get the user
+    let username = data.get("username").unwrap();
+    let user = match User::find_by_mail(username, &conn) {
+        Some(user) => user,
+        None => err!("Username or password is incorrect. Try again.")
+    };
 
-                    device
-                }
-                None => {
-                    // Create new device
-                    Device::new(device_id, user.uuid, device_name, device_type_num)
-                }
+    // Check password
+    let password = data.get("password").unwrap();
+    if !user.check_valid_password(password) {
+        err!("Username or password is incorrect. Try again.")
+    }
+    
+    // Let's only use the header and ignore the 'devicetype' parameter
+    let device_type_num = device_type.0;
+
+    let (device_id, device_name) = match data.is_device {
+        false => { (format!("web-{}", user.uuid), String::from("web")) }
+        true => {
+            (
+                data.get("deviceidentifier").unwrap().clone(),
+                data.get("devicename").unwrap().clone(),
+            )
+        }
+    };
+
+    // Find device or create new
+    let mut device = match Device::find_by_uuid(&device_id, &conn) {
+        Some(device) => {
+            // Check if valid device
+            if device.user_uuid != user.uuid {
+                device.delete(&conn);
+                err!("Device is not owned by user")
             }
+
+            device
+        }
+        None => {
+            // Create new device
+            Device::new(device_id, user.uuid.clone(), device_name, device_type_num)
         }
     };
 
+    let twofactor_token = if user.requires_twofactor() {
+        let twofactor_provider = util::parse_option_string(data.get("twoFactorProvider")).unwrap_or(0);
+        let twofactor_code = match data.get("twoFactorToken") {
+            Some(code) => code,
+            None => err_json!(_json_err_twofactor())
+        };
+
+       match twofactor_provider {
+            0 /* TOTP */ => { 
+                let totp_code: u64 = match twofactor_code.parse() {
+                    Ok(code) => code,
+                    Err(_) => err!("Invalid Totp code")
+                };
+
+                if !user.check_totp_code(totp_code) {
+                    err_json!(_json_err_twofactor())
+                }
+
+                if util::parse_option_string(data.get("twoFactorRemember")).unwrap_or(0) == 1 {
+                    device.refresh_twofactor_remember();
+                    device.twofactor_remember.clone()
+                } else {
+                    device.delete_twofactor_remember();
+                    None
+                }
+            },
+            5 /* Remember */ => {
+                match device.twofactor_remember {
+                    Some(ref remember) if remember == twofactor_code => (),
+                    _ => err_json!(_json_err_twofactor())
+                };
+                None // No twofactor token needed here
+            },
+            _ => err!("Invalid two factor provider"),
+        }
+    } else { None };  // No twofactor token if twofactor is disabled
+
+    // Common
     let user = User::find_by_uuid(&device.user_uuid, &conn).unwrap();
     let orgs = UserOrganization::find_by_user(&user.uuid, &conn);
 
     let (access_token, expires_in) = device.refresh_tokens(&user, orgs);
     device.save(&conn);
 
-    Ok(Json(json!({
+    let mut result = json!({
         "access_token": access_token,
         "expires_in": expires_in,
         "token_type": "Bearer",
         "refresh_token": device.refresh_token,
         "Key": user.key,
-        "PrivateKey": user.private_key
-    })))
+        "PrivateKey": user.private_key,
+        //"TwoFactorToken": "11122233333444555666777888999"
+    });
+
+    if let Some(token) = twofactor_token {
+        result["TwoFactorToken"] = Value::String(token);
+    }
+
+    Ok(Json(result))
+}
+
+fn _json_err_twofactor() -> Value {
+    json!({
+        "error" : "invalid_grant",
+        "error_description" : "Two factor required.",
+        "TwoFactorProviders" : [ 0 ],
+        "TwoFactorProviders2" : { "0" : null }
+    })
+}
+
+/*
+ConnectData {
+    grant_type: Password,
+    is_device: false,
+    data: {
+        "scope": "api offline_access",
+        "client_id": "web",
+        "grant_type": "password",
+        "username": "dani@mail",
+        "password": "8IuV1sJ94tPjyYIK+E+PTjblzjm4W6C4N5wqM0KKsSg="
+    }
+}
+
+RETURNS "TwoFactorToken": "11122233333444555666777888999"
+
+Next login
+ConnectData {
+    grant_type: Password,
+    is_device: false,
+    data: {
+        "scope": "api offline_access",
+        "username": "dani@mail",
+        "client_id": "web",
+        "twofactorprovider": "5",
+        "twofactortoken": "11122233333444555666777888999",
+        "grant_type": "password",
+        "twofactorremember": "0",
+        "password": "8IuV1sJ94tPjyYIK+E+PTjblzjm4W6C4N5wqM0KKsSg="
+    }
 }
+*/
 
 
 struct DeviceType(i32);
diff --git a/src/db/models/device.rs b/src/db/models/device.rs
index e3f1d53f..ee3f595c 100644
--- a/src/db/models/device.rs
+++ b/src/db/models/device.rs
@@ -19,6 +19,8 @@ pub struct Device {
     pub push_token: Option<String>,
 
     pub refresh_token: String,
+
+    pub twofactor_remember: Option<String>,
 }
 
 /// Local methods
@@ -37,9 +39,22 @@ impl Device {
 
             push_token: None,
             refresh_token: String::new(),
+            twofactor_remember: None,
         }
     }
 
+    pub fn refresh_twofactor_remember(&mut self) {
+        use data_encoding::BASE64;
+        use crypto;
+
+        self.twofactor_remember = Some(BASE64.encode(&crypto::get_random(vec![0u8; 180])));
+    }
+
+    pub fn delete_twofactor_remember(&mut self) {
+        self.twofactor_remember = None;
+    }
+
+
     pub fn refresh_tokens(&mut self, user: &super::User, orgs: Vec<super::UserOrganization>) -> (String, i64) {
         // If there is no refresh token, we create one
         if self.refresh_token.is_empty() {
diff --git a/src/db/models/user.rs b/src/db/models/user.rs
index 0287d459..891088b4 100644
--- a/src/db/models/user.rs
+++ b/src/db/models/user.rs
@@ -26,8 +26,10 @@ pub struct User {
     pub key: String,
     pub private_key: Option<String>,
     pub public_key: Option<String>,
+    
     pub totp_secret: Option<String>,
     pub totp_recover: Option<String>,
+
     pub security_stamp: String,
 
     pub equivalent_domains: String,
@@ -61,6 +63,7 @@ impl User {
             password_hint: None,
             private_key: None,
             public_key: None,
+            
             totp_secret: None,
             totp_recover: None,
 
@@ -95,23 +98,23 @@ impl User {
         self.security_stamp = Uuid::new_v4().to_string();
     }
 
-    pub fn check_totp_code(&self, totp_code: Option<u64>) -> bool {
+    pub fn requires_twofactor(&self) -> bool {
+        self.totp_secret.is_some()
+    }
+
+    pub fn check_totp_code(&self, totp_code: u64) -> bool {
         if let Some(ref totp_secret) = self.totp_secret {
-            if let Some(code) = totp_code {
-                // Validate totp
-                use data_encoding::BASE32;
-                use oath::{totp_raw_now, HashType};
-
-                let decoded_secret = match BASE32.decode(totp_secret.as_bytes()) {
-                    Ok(s) => s,
-                    Err(_) => return false
-                };
-
-                let generated = totp_raw_now(&decoded_secret, 6, 0, 30, &HashType::SHA1);
-                generated == code
-            } else {
-                false
-            }
+            // Validate totp
+            use data_encoding::BASE32;
+            use oath::{totp_raw_now, HashType};
+
+            let decoded_secret = match BASE32.decode(totp_secret.as_bytes()) {
+                Ok(s) => s,
+                Err(_) => return false
+            };
+
+            let generated = totp_raw_now(&decoded_secret, 6, 0, 30, &HashType::SHA1);
+            generated == totp_code
         } else {
             true
         }
diff --git a/src/db/schema.rs b/src/db/schema.rs
index c144e9b8..baf9e6d2 100644
--- a/src/db/schema.rs
+++ b/src/db/schema.rs
@@ -24,6 +24,13 @@ table! {
     }
 }
 
+table! {
+    ciphers_collections (cipher_uuid, collection_uuid) {
+        cipher_uuid -> Text,
+        collection_uuid -> Text,
+    }
+}
+
 table! {
     collections (uuid) {
         uuid -> Text,
@@ -43,6 +50,7 @@ table! {
         type_ -> Integer,
         push_token -> Nullable<Text>,
         refresh_token -> Text,
+        twofactor_remember -> Nullable<Text>,
     }
 }
 
@@ -101,13 +109,6 @@ table! {
     }
 }
 
-table! {
-    ciphers_collections (cipher_uuid, collection_uuid) {
-        cipher_uuid -> Text,
-        collection_uuid -> Text,
-    }
-}
-
 table! {
     users_organizations (uuid) {
         uuid -> Text,
@@ -124,6 +125,8 @@ table! {
 joinable!(attachments -> ciphers (cipher_uuid));
 joinable!(ciphers -> organizations (organization_uuid));
 joinable!(ciphers -> users (user_uuid));
+joinable!(ciphers_collections -> ciphers (cipher_uuid));
+joinable!(ciphers_collections -> collections (collection_uuid));
 joinable!(collections -> organizations (org_uuid));
 joinable!(devices -> users (user_uuid));
 joinable!(folders -> users (user_uuid));
@@ -131,14 +134,13 @@ joinable!(folders_ciphers -> ciphers (cipher_uuid));
 joinable!(folders_ciphers -> folders (folder_uuid));
 joinable!(users_collections -> collections (collection_uuid));
 joinable!(users_collections -> users (user_uuid));
-joinable!(ciphers_collections -> collections (collection_uuid));
-joinable!(ciphers_collections -> ciphers (cipher_uuid));
 joinable!(users_organizations -> organizations (org_uuid));
 joinable!(users_organizations -> users (user_uuid));
 
 allow_tables_to_appear_in_same_query!(
     attachments,
     ciphers,
+    ciphers_collections,
     collections,
     devices,
     folders,
@@ -146,6 +148,5 @@ allow_tables_to_appear_in_same_query!(
     organizations,
     users,
     users_collections,
-    ciphers_collections,
     users_organizations,
 );
-- 
GitLab