diff --git a/src/error.rs b/src/error.rs
index 18e3bee..0f7a34c 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -75,6 +75,12 @@ impl LuksError {
     }
 }
 
+impl From<LuksError> for Fido2LuksError {
+    fn from(e: LuksError) -> Self {
+        Fido2LuksError::LuksError { cause: e }
+    }
+}
+
 use libcryptsetup_rs::LibcryptErr;
 use std::io::ErrorKind;
 use std::string::FromUtf8Error;
diff --git a/src/luks.rs b/src/luks.rs
index c38359c..efdc0fc 100644
--- a/src/luks.rs
+++ b/src/luks.rs
@@ -7,18 +7,240 @@ use libcryptsetup_rs::{
 use std::collections::{HashMap, HashSet};
 use std::path::Path;
 
-fn load_device_handle<P: AsRef<Path>>(path: P) -> Fido2LuksResult<CryptDevice> {
-    let mut device = CryptInit::init(path.as_ref())?;
-    device.context_handle().load::<()>(None, None)?;
-    Ok(device)
+struct LuksDevice {
+    device: CryptDevice,
+    luks2: Option<bool>,
 }
 
-fn check_luks2(device: &mut CryptDevice) -> Fido2LuksResult<()> {
-    match device.format_handle().get_type()? {
-        EncryptionFormat::Luks2 => Ok(()),
-        _ => Err(Fido2LuksError::LuksError {
-            cause: LuksError::Luks2Required,
-        }),
+impl LuksDevice {
+    pub fn load<P: AsRef<Path>>(path: P) -> Fido2LuksResult<LuksDevice> {
+        let mut device = CryptInit::init(path.as_ref())?;
+        device.context_handle().load::<()>(None, None)?;
+        Ok(Self {
+            device,
+            luks2: None,
+        })
+    }
+
+    pub fn is_luks2(&mut self) -> Fido2LuksResult<bool> {
+        if let Some(luks2) = self.luks2 {
+            Ok(luks2)
+        } else {
+            self.luks2 = Some(match self.device.format_handle().get_type()? {
+                EncryptionFormat::Luks2 => true,
+                _ => false,
+            });
+            self.is_luks2()
+        }
+    }
+
+    pub fn tokens<'a>(
+        &'a mut self,
+    ) -> Fido2LuksResult<Box<dyn Iterator<Item = Fido2LuksResult<(u32, Fido2LuksToken)>> + 'a>>
+    {
+        Ok(Box::new((0..32).filter_map(move |i| {
+            let status = match self.device.token_handle().status(i) {
+                Ok(status) => status,
+                Err(err) => return Some(Err(Fido2LuksError::from(err))),
+            };
+            match status {
+                CryptTokenInfo::Inactive => return None,
+                CryptTokenInfo::Internal(s)
+                | CryptTokenInfo::InternalUnknown(s)
+                | CryptTokenInfo::ExternalUnknown(s)
+                | CryptTokenInfo::External(s)
+                    if &s != "fido2luks" =>
+                {
+                    return None
+                }
+                _ => (),
+            };
+            let json = match self.device.token_handle().json_get(i) {
+                Ok(json) => json,
+                Err(err) => return Some(Err(Fido2LuksError::from(err))),
+            };
+            let info: Fido2LuksToken =
+                match serde_json::from_value(json.clone()).map_err(|_| Fido2LuksError::LuksError {
+                    cause: LuksError::InvalidToken(json.to_string()),
+                }) {
+                    Ok(info) => info,
+                    Err(err) => return Some(Err(Fido2LuksError::from(err))),
+                };
+            Some(Ok((i, info)))
+        })))
+    }
+
+    pub fn find_token(&mut self, slot: u32) -> Fido2LuksResult<Option<(u32, Fido2LuksToken)>> {
+        let slot_str = slot.to_string();
+        for token in self.tokens()? {
+            let (id, token) = token?;
+            if token.keyslots.contains(&slot_str) {
+                return Ok(Some((id, token)));
+            }
+        }
+        Ok(None)
+    }
+
+    pub fn add_key(
+        &mut self,
+        secret: &[u8],
+        old_secret: &[u8],
+        iteration_time: Option<u64>,
+        credential_id: Option<&[u8]>,
+    ) -> Fido2LuksResult<u32> {
+        if let Some(millis) = iteration_time {
+            self.device.settings_handle().set_iteration_time(millis)
+        }
+        let slot = self
+            .device
+            .keyslot_handle()
+            .add_by_passphrase(None, old_secret, secret)?;
+        if let Some(id) = credential_id {
+            self.device.token_handle().json_set(TokenInput::AddToken(
+                &serde_json::to_value(&Fido2LuksToken::new(id, slot)).unwrap(),
+            ))?;
+        }
+
+        Ok(slot)
+    }
+
+    pub fn remove_keyslots(&mut self, exclude: &[u32]) -> Fido2LuksResult<u32> {
+        let mut destroyed = 0;
+        let mut tokens = Vec::new();
+        for slot in 0..256 {
+            match self.device.keyslot_handle().status(slot)? {
+                KeyslotInfo::Inactive => continue,
+                KeyslotInfo::Active | KeyslotInfo::ActiveLast if !exclude.contains(&slot) => {
+                    if self.is_luks2()? {
+                        if let Some((id, _token)) = self.find_token(slot)? {
+                            tokens.push(id);
+                        }
+                    }
+                    self.device.keyslot_handle().destroy(slot)?;
+                    destroyed += 1;
+                }
+                KeyslotInfo::ActiveLast => break,
+                _ => (),
+            }
+            if self.device.keyslot_handle().status(slot)? == KeyslotInfo::ActiveLast {
+                break;
+            }
+        }
+        // Ensure indices stay valid
+        tokens.sort();
+        for token in tokens.iter().rev() {
+            self.device
+                .token_handle()
+                .json_set(TokenInput::RemoveToken(*token))?;
+        }
+        Ok(destroyed)
+    }
+
+    pub fn replace_key(
+        &mut self,
+        secret: &[u8],
+        old_secret: &[u8],
+        iteration_time: Option<u64>,
+        credential_id: Option<&[u8]>,
+    ) -> Fido2LuksResult<u32> {
+        if let Some(millis) = iteration_time {
+            self.device.settings_handle().set_iteration_time(millis)
+        }
+        // Use activate dry-run to locate keyslot
+        let slot = self.device.activate_handle().activate_by_passphrase(
+            None,
+            None,
+            old_secret,
+            CryptActivateFlags::empty(),
+        )?;
+        self.device.keyslot_handle().change_by_passphrase(
+            Some(slot),
+            Some(slot),
+            old_secret,
+            secret,
+        )? as u32;
+        if let Some(id) = credential_id {
+            if self.is_luks2()? {
+                let token = self.find_token(slot)?.map(|(t, _)| t);
+                let json = serde_json::to_value(&Fido2LuksToken::new(id, slot)).unwrap();
+                if let Some(token) = token {
+                    self.device
+                        .token_handle()
+                        .json_set(TokenInput::ReplaceToken(token, &json))?;
+                } else {
+                    self.device
+                        .token_handle()
+                        .json_set(TokenInput::AddToken(&json))?;
+                }
+            }
+        }
+        Ok(slot)
+    }
+
+    pub fn activate(
+        &mut self,
+        name: &str,
+        secret: &[u8],
+        slot_hint: Option<u32>,
+    ) -> Fido2LuksResult<u32> {
+        self.device
+            .activate_handle()
+            .activate_by_passphrase(Some(name), slot_hint, secret, CryptActivateFlags::empty())
+            .map_err(|_e| Fido2LuksError::WrongSecret)
+    }
+
+    pub fn activate_token(
+        &mut self,
+        name: &str,
+        secret: impl Fn(Vec<String>) -> Fido2LuksResult<([u8; 32], String)>,
+        slot_hint: Option<u32>,
+    ) -> Fido2LuksResult<u32> {
+        if !self.is_luks2()? {
+            return Err(LuksError::Luks2Required.into());
+        }
+        let mut creds: HashMap<String, HashSet<u32>> = HashMap::new();
+        for token in self.tokens()? {
+            let token = match token {
+                Ok((_id, t)) => t,
+                _ => continue, // An corrupted token should't lock the user out
+            };
+            let slots = || {
+                token
+                    .keyslots
+                    .iter()
+                    .filter_map(|slot| slot.parse::<u32>().ok())
+            };
+            for cred in token.credential.iter() {
+                creds
+                    .entry(cred.clone())
+                    .or_insert_with(|| slots().collect::<HashSet<u32>>())
+                    .extend(slots());
+            }
+        }
+        if creds.is_empty() {
+            return Err(Fido2LuksError::LuksError {
+                cause: LuksError::NoToken,
+            });
+        }
+        let (secret, credential) = secret(creds.keys().cloned().collect())?;
+        let slots = creds.get(&credential).unwrap();
+        let slots = slots
+            .iter()
+            .cloned()
+            .map(Option::Some)
+            .chain(std::iter::once(None).take(slots.is_empty() as usize));
+        for slot in slots {
+            match self
+                .device
+                .activate_handle()
+                .activate_by_passphrase(Some(name), slot, &secret, CryptActivateFlags::empty())
+                .map_err(LuksError::activate)
+            {
+                Err(Fido2LuksError::WrongSecret) => (),
+                res => return res,
+            }
+        }
+        self.activate(name, &secret, slot_hint)
     }
 }
 
@@ -46,12 +268,9 @@ pub fn open_container<P: AsRef<Path>>(
     secret: &[u8],
     slot_hint: Option<u32>,
 ) -> Fido2LuksResult<()> {
-    let mut device = load_device_handle(path)?;
-    device
-        .activate_handle()
-        .activate_by_passphrase(Some(name), slot_hint, secret, CryptActivateFlags::empty())
-        .map(|_slot| ())
-        .map_err(|_e| Fido2LuksError::WrongSecret)
+    LuksDevice::load(path)?
+        .activate(name, secret, slot_hint)
+        .map(|_| ())
 }
 
 pub fn open_container_token<P: AsRef<Path>>(
@@ -59,65 +278,9 @@ pub fn open_container_token<P: AsRef<Path>>(
     name: &str,
     secret: impl Fn(Vec<String>) -> Fido2LuksResult<([u8; 32], String)>,
 ) -> Fido2LuksResult<()> {
-    let mut device = load_device_handle(path)?;
-    check_luks2(&mut device)?;
-
-    let mut creds = HashMap::new();
-    for i in 0..256 {
-        let status = device.token_handle().status(i)?;
-        match status {
-            CryptTokenInfo::Inactive => break,
-            CryptTokenInfo::Internal(s)
-            | CryptTokenInfo::InternalUnknown(s)
-            | CryptTokenInfo::ExternalUnknown(s)
-            | CryptTokenInfo::External(s)
-                if &s != "fido2luks" =>
-            {
-                continue
-            }
-            _ => (),
-        };
-        let json = device.token_handle().json_get(i)?;
-        let info: Fido2LuksToken =
-            serde_json::from_value(json.clone()).map_err(|_| Fido2LuksError::LuksError {
-                cause: LuksError::InvalidToken(json.to_string()),
-            })?;
-        let slots = || {
-            info.keyslots
-                .iter()
-                .filter_map(|slot| slot.parse::<u32>().ok())
-        };
-        for cred in info.credential.iter().cloned() {
-            creds
-                .entry(cred)
-                .or_insert_with(|| slots().collect::<HashSet<u32>>())
-                .extend(slots());
-        }
-    }
-    if creds.is_empty() {
-        return Err(Fido2LuksError::LuksError {
-            cause: LuksError::NoToken,
-        });
-    }
-    let (secret, credential) = secret(creds.keys().cloned().collect())?;
-    let slots = creds.get(&credential).unwrap();
-    let slots = slots
-        .iter()
-        .cloned()
-        .map(Option::Some)
-        .chain(std::iter::once(None).take(slots.is_empty() as usize));
-    for slot in slots {
-        match device
-            .activate_handle()
-            .activate_by_passphrase(Some(name), slot, &secret, CryptActivateFlags::empty())
-            .map(|_slot| ())
-            .map_err(LuksError::activate)
-        {
-            Err(Fido2LuksError::WrongSecret) => (),
-            res => return res,
-        }
-    }
-    Err(Fido2LuksError::WrongSecret)
+    LuksDevice::load(path)?
+        .activate_token(name, secret, None)
+        .map(|_| ())
 }
 
 pub fn add_key<P: AsRef<Path>>(
@@ -127,88 +290,11 @@ pub fn add_key<P: AsRef<Path>>(
     iteration_time: Option<u64>,
     credential_id: Option<&[u8]>,
 ) -> Fido2LuksResult<u32> {
-    let mut device = load_device_handle(path)?;
-    if let Some(millis) = iteration_time {
-        device.settings_handle().set_iteration_time(millis)
-    }
-    let slot = device
-        .keyslot_handle()
-        .add_by_passphrase(None, old_secret, secret)?;
-    if let Some(id) = credential_id {
-        /*  if let e @ Err(_) = check_luks2(&mut device) {
-            //rollback
-            device.keyslot_handle(Some(slot)).destroy()?;
-            return e.map(|_| 0u32);
-        }*/
-        device.token_handle().json_set(TokenInput::AddToken(
-            &serde_json::to_value(&Fido2LuksToken::new(id, slot)).unwrap(),
-        ))?;
-    }
-
-    Ok(slot)
-}
-
-fn find_token(
-    device: &mut CryptDevice,
-    slot: u32,
-) -> Fido2LuksResult<Option<(u32, Fido2LuksToken)>> {
-    for i in 0..256 {
-        let status = device.token_handle().status(i)?;
-        match status {
-            CryptTokenInfo::Inactive => break,
-            CryptTokenInfo::Internal(s)
-            | CryptTokenInfo::InternalUnknown(s)
-            | CryptTokenInfo::ExternalUnknown(s)
-            | CryptTokenInfo::External(s)
-                if &s != "fido2luks" =>
-            {
-                continue
-            }
-            _ => (),
-        };
-        let json = device.token_handle().json_get(i)?;
-        let info: Fido2LuksToken =
-            serde_json::from_value(json.clone()).map_err(|_| Fido2LuksError::LuksError {
-                cause: LuksError::InvalidToken(json.to_string()),
-            })?;
-        if info.keyslots.contains(&slot.to_string()) {
-            return Ok(Some((i, info)));
-        }
-    }
-    Ok(None)
+    LuksDevice::load(path)?.add_key(secret, old_secret, iteration_time, credential_id)
 }
 
 pub fn remove_keyslots<P: AsRef<Path>>(path: P, exclude: &[u32]) -> Fido2LuksResult<u32> {
-    let mut device = load_device_handle(path)?;
-    let mut destroyed = 0;
-    let mut tokens = Vec::new();
-    for slot in 0..256 {
-        match device.keyslot_handle().status(slot)? {
-            KeyslotInfo::Inactive => continue,
-            KeyslotInfo::Active | KeyslotInfo::ActiveLast if !exclude.contains(&slot) => {
-                if let Ok(_) = check_luks2(&mut device) {
-                    if let Some((token, _)) = find_token(&mut device, slot)? {
-                        tokens.push(token);
-                    }
-                }
-                device.keyslot_handle().destroy(slot)?;
-                destroyed += 1;
-            }
-            KeyslotInfo::ActiveLast => break,
-            _ => (),
-        }
-        if device.keyslot_handle().status(slot)? == KeyslotInfo::ActiveLast {
-            break;
-        }
-    }
-    // Ensure indices stay valid
-    tokens.sort();
-    for token in tokens.iter().rev() {
-        device
-            .token_handle()
-            .json_set(TokenInput::RemoveToken(*token))?;
-    }
-    Ok(destroyed)
+    LuksDevice::load(path)?.remove_keyslots(exclude)
 }
 
 pub fn replace_key<P: AsRef<Path>>(
@@ -218,34 +304,5 @@ pub fn replace_key<P: AsRef<Path>>(
     iteration_time: Option<u64>,
     credential_id: Option<&[u8]>,
 ) -> Fido2LuksResult<u32> {
-    let mut device = load_device_handle(path)?;
-    if let Some(millis) = iteration_time {
-        device.settings_handle().set_iteration_time(millis)
-    }
-    // Use activate dry-run to locate keyslot
-    let slot = device.activate_handle().activate_by_passphrase(
-        None,
-        None,
-        old_secret,
-        CryptActivateFlags::empty(),
-    )?;
-    device
-        .keyslot_handle()
-        .change_by_passphrase(Some(slot), Some(slot), old_secret, secret)? as u32;
-    if let Some(id) = credential_id {
-        if check_luks2(&mut device).is_ok() {
-            let token = find_token(&mut device, slot)?.map(|(t, _)| t);
-            let json = serde_json::to_value(&Fido2LuksToken::new(id, slot)).unwrap();
-            if let Some(token) = token {
-                device
-                    .token_handle()
-                    .json_set(TokenInput::ReplaceToken(token, &json))?;
-            } else {
-                device
-                    .token_handle()
-                    .json_set(TokenInput::AddToken(&json))?;
-            }
-        }
-    }
-    Ok(slot)
+    LuksDevice::load(path)?.replace_key(secret, old_secret, iteration_time, credential_id)
 }