// Copyright 2021 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Wireless network security descriptors and authenticators.
//!
//! This module describes wireless network security protocols as well as credentials used to
//! authenticate using those protocols. Types in this module provide two important features:
//! conversions from and into FIDL types used within the WLAN stack and consistency. Here,
//! consistency means that types have no invalid values; given an instance of a type, it is always
//! valid. For example, a `SecurityAuthenticator` always represents a valid protocol-credential
//! pair.
//!
//! The `SecurityDescriptor` and `SecurityAuthenticator` types form the primary API of this module.
//! A _descriptor_ merely describes (names) a security protocol such as WPA2 Personal using TKIP.
//! An _authenticator_ describes a security protocol and additionally contains credentials used to
//! authenticate using that protocol such as WPA Personal using TKIP with a PSK credential.
//! Authenticators can be converted into descriptors (which drops any associated credentials).
//!
//! # FIDL
//!
//! Types in this module support conversion into and from datagrams in the
//! `fuchsia.wlan.common.security` package. When interacting with FIDL, most code should prefer
//! conversions between FIDL types and this module's `SecurityDescriptor` and
//! `SecurityAuthenticator` types, though conversions for intermediate types are also provided.
//!
//! The models used by this module and the FIDL package differ, so types do not always have a
//! direct analog, but the table below describes the most straightforward and important mappings.
//!
//! | Rust Type               | FIDL Type        | Rust to FIDL | FIDL to Rust |
//! |-------------------------|------------------|--------------|--------------|
//! | `SecurityDescriptor`    | `Protocol`       | `From`       | `From`       |
//! | `SecurityAuthenticator` | `Authentication` | `From`       | `TryFrom`    |
//!
//! # Cryptography and RSN
//!
//! This module does **not** implement any security protocols nor cryptography and only describes
//! them with limited detail as needed at API boundaries and in code that merely interacts with
//! these protocols. See the `rsn` crate for implementations of the IEEE Std 802.11-2016 Robust
//! Security Network (RSN).

// NOTE: At the time of this writing, the WLAN stack does not support WPA Enterprise. One
//       consequence of this is that conversions between Rust and FIDL types may return errors or
//       panic when they involve WPA Enterprise representations. See https://fxbug.dev/42174395 and the TODOs
//       in this module.

pub mod wep;
pub mod wpa;

use fidl_fuchsia_wlan_common_security as fidl_security;
use thiserror::Error;

use crate::security::{
    wep::WepKey,
    wpa::credential::{Passphrase, PassphraseError, Psk, PskError},
};

/// Extension methods for the `Credentials` FIDL datagram.
pub trait CredentialsExt {
    fn into_wep(self) -> Option<fidl_security::WepCredentials>;
    fn into_wpa(self) -> Option<fidl_security::WpaCredentials>;
}

impl CredentialsExt for fidl_security::Credentials {
    fn into_wep(self) -> Option<fidl_security::WepCredentials> {
        if let fidl_security::Credentials::Wep(credentials) = self {
            Some(credentials)
        } else {
            None
        }
    }

    fn into_wpa(self) -> Option<fidl_security::WpaCredentials> {
        if let fidl_security::Credentials::Wpa(credentials) = self {
            Some(credentials)
        } else {
            None
        }
    }
}

#[derive(Clone, Copy, Debug, Error, Eq, PartialEq)]
#[non_exhaustive]
pub enum SecurityError {
    #[error(transparent)]
    Wep(#[from] wep::WepError),
    #[error(transparent)]
    Wpa(#[from] wpa::WpaError),
    /// This error occurs when there is an incompatibility between security protocols, features,
    /// and/or credentials.
    ///
    /// Note that this is distinct from `SecurityError::Unsupported`.
    #[error("incompatible protocol or features")]
    Incompatible,
    /// This error occurs when a specified security protocol, features, and/or credentials are
    /// **not** supported.
    ///
    /// Note that this is distinct from `SecurityError::Incompatible`. Unsupported features may be
    /// compatible and specified in IEEE Std. 802.11-2016.
    #[error("unsupported protocol or features")]
    Unsupported,
}

impl From<PassphraseError> for SecurityError {
    fn from(error: PassphraseError) -> Self {
        SecurityError::Wpa(error.into())
    }
}

impl From<PskError> for SecurityError {
    fn from(error: PskError) -> Self {
        SecurityError::Wpa(error.into())
    }
}

/// General credential data that is not explicitly coupled to a particular security protocol.
///
/// The variants of this enumeration are particular to general protocols (i.e., WEP and WPA), but
/// don't provide any more details or validation. For WPA credential data, this means that the
/// version of the WPA security protocol is entirely unknown.
///
/// This type is meant for code and APIs that accept such bare credentials and must incorporate
/// additional information or apply heuristics to negotiate a specific protocol. For example, this
/// occurs in code that communicates directly with SME without support from the Policy layer to
/// derive this information.
///
/// The FIDL analogue of this type is `fuchsia.wlan.common.security.Credentials`, into and from
/// which this type can be infallibly converted.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum BareCredentials {
    /// WEP key.
    WepKey(WepKey),
    /// WPA passphrase.
    ///
    /// Passphrases can be used to authenticate with WPA1, WPA2, and WPA3.
    WpaPassphrase(Passphrase),
    /// WPA PSK.
    ///
    /// PSKs are distinct from passphrases and can be used to authenticate with WPA1 and WPA2. A
    /// PSK cannot be used to authenticate with WPA3.
    WpaPsk(Psk),
}

impl From<BareCredentials> for fidl_security::Credentials {
    fn from(credentials: BareCredentials) -> Self {
        match credentials {
            BareCredentials::WepKey(key) => {
                fidl_security::Credentials::Wep(fidl_security::WepCredentials { key: key.into() })
            }
            BareCredentials::WpaPassphrase(passphrase) => fidl_security::Credentials::Wpa(
                fidl_security::WpaCredentials::Passphrase(passphrase.into()),
            ),
            BareCredentials::WpaPsk(psk) => {
                fidl_security::Credentials::Wpa(fidl_security::WpaCredentials::Psk(psk.into()))
            }
        }
    }
}

impl TryFrom<fidl_security::Credentials> for BareCredentials {
    type Error = SecurityError;

    fn try_from(credentials: fidl_security::Credentials) -> Result<Self, Self::Error> {
        match credentials {
            fidl_security::Credentials::Wep(fidl_security::WepCredentials { key }) => {
                WepKey::try_from_literal_bytes(key.as_slice())
                    .map(|key| BareCredentials::WepKey(key))
                    .map_err(From::from)
            }
            fidl_security::Credentials::Wpa(credentials) => match credentials {
                fidl_security::WpaCredentials::Passphrase(passphrase) => {
                    Passphrase::try_from(passphrase)
                        .map(|passphrase| BareCredentials::WpaPassphrase(passphrase))
                        .map_err(From::from)
                }
                fidl_security::WpaCredentials::Psk(psk) => Psk::try_from(psk.as_slice())
                    .map(|psk| BareCredentials::WpaPsk(psk))
                    .map_err(From::from),
                // Unknown variant.
                _ => Err(SecurityError::Incompatible),
            },
            // Unknown variant.
            _ => Err(SecurityError::Incompatible),
        }
    }
}

/// Conversion from a WPA passphrase into bare credentials.
impl From<Passphrase> for BareCredentials {
    fn from(passphrase: Passphrase) -> Self {
        BareCredentials::WpaPassphrase(passphrase)
    }
}

/// Conversion from a WPA PSK into bare credentials.
impl From<Psk> for BareCredentials {
    fn from(psk: Psk) -> Self {
        BareCredentials::WpaPsk(psk)
    }
}

/// Conversion from a WEP key into bare credentials.
impl From<WepKey> for BareCredentials {
    fn from(key: WepKey) -> Self {
        BareCredentials::WepKey(key)
    }
}

/// Description of a wireless network security protocol.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum SecurityDescriptor {
    Open,
    Wep,
    Wpa(wpa::WpaDescriptor),
}

impl SecurityDescriptor {
    /// Open (no user authentication nor traffic encryption).
    pub const OPEN: Self = SecurityDescriptor::Open;
    /// WEP (trivially insecure; for legacy support only).
    ///
    /// This protocol is not configurable beyond the format of credentials used to authenticate.
    /// WEP provides no protection and is provided for legacy support only.
    pub const WEP: Self = SecurityDescriptor::Wep;
    /// Legacy WPA (WPA1).
    ///
    /// This protocol is not configurable beyond the format of credentials used to authenticate.
    pub const WPA1: Self = SecurityDescriptor::Wpa(wpa::WpaDescriptor::Wpa1 { credentials: () });
    /// WPA2 Personal.
    ///
    /// Describes the personal variant of the WPA2 protocol. This descriptor does not specify a
    /// pairwise cipher.
    pub const WPA2_PERSONAL: Self = SecurityDescriptor::Wpa(wpa::WpaDescriptor::Wpa2 {
        cipher: None,
        authentication: wpa::Authentication::Personal(()),
    });
    /// WPA3 Personal.
    ///
    /// Describes the personal variant of the WPA3 protocol. This descriptor does not specify a
    /// pairwise cipher.
    pub const WPA3_PERSONAL: Self = SecurityDescriptor::Wpa(wpa::WpaDescriptor::Wpa3 {
        cipher: None,
        authentication: wpa::Authentication::Personal(()),
    });

    /// Binds bare credentials to a descriptor to form an authenticator.
    ///
    /// A security descriptor only describes a protocol and bare credentials provide authentication
    /// data without completely describing a protocol. When compatible, a descriptor and
    /// credentials form the components of an authenticator, and this function attempts to form an
    /// authenticator by binding these components together.
    ///
    /// # Errors
    ///
    /// Returns an error if the bare credentials are incompatible with the descriptor.
    pub fn bind(
        self,
        credentials: Option<BareCredentials>,
    ) -> Result<SecurityAuthenticator, SecurityError> {
        match self {
            SecurityDescriptor::Open if credentials.is_none() => Ok(SecurityAuthenticator::Open),
            SecurityDescriptor::Wep => match credentials {
                Some(BareCredentials::WepKey(key)) => {
                    Ok(SecurityAuthenticator::Wep(wep::WepAuthenticator { key }))
                }
                _ => Err(SecurityError::Incompatible),
            },
            SecurityDescriptor::Wpa(wpa) => match credentials {
                Some(credentials) => wpa.bind(credentials).map(SecurityAuthenticator::Wpa),
                _ => Err(SecurityError::Incompatible),
            },
            _ => Err(SecurityError::Incompatible),
        }
    }

    pub fn is_open(&self) -> bool {
        matches!(self, SecurityDescriptor::Open)
    }

    pub fn is_wep(&self) -> bool {
        matches!(self, SecurityDescriptor::Wep)
    }

    pub fn is_wpa(&self) -> bool {
        matches!(self, SecurityDescriptor::Wpa(_))
    }
}

impl From<fidl_security::Protocol> for SecurityDescriptor {
    fn from(protocol: fidl_security::Protocol) -> Self {
        match protocol {
            fidl_security::Protocol::Open => SecurityDescriptor::Open,
            fidl_security::Protocol::Wep => SecurityDescriptor::Wep,
            fidl_security::Protocol::Wpa1 => {
                SecurityDescriptor::Wpa(wpa::WpaDescriptor::Wpa1 { credentials: () })
            }
            fidl_security::Protocol::Wpa2Personal => {
                SecurityDescriptor::Wpa(wpa::WpaDescriptor::Wpa2 {
                    authentication: wpa::Authentication::Personal(()),
                    cipher: None,
                })
            }
            fidl_security::Protocol::Wpa2Enterprise => {
                SecurityDescriptor::Wpa(wpa::WpaDescriptor::Wpa2 {
                    authentication: wpa::Authentication::Enterprise(()),
                    cipher: None,
                })
            }
            fidl_security::Protocol::Wpa3Personal => {
                SecurityDescriptor::Wpa(wpa::WpaDescriptor::Wpa3 {
                    authentication: wpa::Authentication::Personal(()),
                    cipher: None,
                })
            }
            fidl_security::Protocol::Wpa3Enterprise => {
                SecurityDescriptor::Wpa(wpa::WpaDescriptor::Wpa3 {
                    authentication: wpa::Authentication::Enterprise(()),
                    cipher: None,
                })
            }
            _ => panic!("unknown FIDL security protocol variant"),
        }
    }
}

impl From<SecurityDescriptor> for fidl_security::Protocol {
    fn from(descriptor: SecurityDescriptor) -> Self {
        match descriptor {
            SecurityDescriptor::Open => fidl_security::Protocol::Open,
            SecurityDescriptor::Wep => fidl_security::Protocol::Wep,
            SecurityDescriptor::Wpa(wpa) => match wpa {
                wpa::WpaDescriptor::Wpa1 { .. } => fidl_security::Protocol::Wpa1,
                wpa::WpaDescriptor::Wpa2 { authentication, .. } => match authentication {
                    wpa::Authentication::Personal(_) => fidl_security::Protocol::Wpa2Personal,
                    wpa::Authentication::Enterprise(_) => fidl_security::Protocol::Wpa2Enterprise,
                },
                wpa::WpaDescriptor::Wpa3 { authentication, .. } => match authentication {
                    wpa::Authentication::Personal(_) => fidl_security::Protocol::Wpa3Personal,
                    wpa::Authentication::Enterprise(_) => fidl_security::Protocol::Wpa3Enterprise,
                },
            },
        }
    }
}

impl From<wpa::WpaDescriptor> for SecurityDescriptor {
    fn from(descriptor: wpa::WpaDescriptor) -> Self {
        SecurityDescriptor::Wpa(descriptor)
    }
}

/// Credentials and configuration for authenticating using a particular wireless network security
/// protocol.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum SecurityAuthenticator {
    Open,
    Wep(wep::WepAuthenticator),
    Wpa(wpa::WpaAuthenticator),
}

impl SecurityAuthenticator {
    /// Converts an authenticator into a descriptor with no payload (the unit type `()`). Any
    /// payload (i.e., credentials) are dropped.
    pub fn into_descriptor(self) -> SecurityDescriptor {
        match self {
            SecurityAuthenticator::Open => SecurityDescriptor::Open,
            SecurityAuthenticator::Wep(_) => SecurityDescriptor::Wep,
            SecurityAuthenticator::Wpa(authenticator) => {
                SecurityDescriptor::Wpa(authenticator.into_descriptor())
            }
        }
    }

    pub fn into_wep(self) -> Option<wep::WepAuthenticator> {
        match self {
            SecurityAuthenticator::Wep(authenticator) => Some(authenticator),
            _ => None,
        }
    }

    pub fn into_wpa(self) -> Option<wpa::WpaAuthenticator> {
        match self {
            SecurityAuthenticator::Wpa(authenticator) => Some(authenticator),
            _ => None,
        }
    }

    /// Converts the authenticator to bare credentials, if any.
    ///
    /// Returns `None` if the authenticator is `SecurityAuthenticator::None`, as there are no
    /// corresponding credentials in this case.
    pub fn to_credentials(&self) -> Option<BareCredentials> {
        match self {
            SecurityAuthenticator::Open => None,
            SecurityAuthenticator::Wep(wep::WepAuthenticator { ref key }) => {
                Some(key.clone().into())
            }
            SecurityAuthenticator::Wpa(ref wpa) => Some(wpa.to_credentials().into()),
        }
    }

    pub fn as_wep(&self) -> Option<&wep::WepAuthenticator> {
        match self {
            SecurityAuthenticator::Wep(ref authenticator) => Some(authenticator),
            _ => None,
        }
    }

    pub fn as_wpa(&self) -> Option<&wpa::WpaAuthenticator> {
        match self {
            SecurityAuthenticator::Wpa(ref authenticator) => Some(authenticator),
            _ => None,
        }
    }

    pub fn is_open(&self) -> bool {
        matches!(self, SecurityAuthenticator::Open)
    }

    pub fn is_wep(&self) -> bool {
        matches!(self, SecurityAuthenticator::Wep(_))
    }

    pub fn is_wpa(&self) -> bool {
        matches!(self, SecurityAuthenticator::Wpa(_))
    }
}

impl From<wep::WepAuthenticator> for SecurityAuthenticator {
    fn from(authenticator: wep::WepAuthenticator) -> Self {
        SecurityAuthenticator::Wep(authenticator)
    }
}

impl From<wpa::WpaAuthenticator> for SecurityAuthenticator {
    fn from(authenticator: wpa::WpaAuthenticator) -> Self {
        SecurityAuthenticator::Wpa(authenticator)
    }
}

impl From<SecurityAuthenticator> for fidl_security::Authentication {
    fn from(authenticator: SecurityAuthenticator) -> Self {
        match authenticator {
            SecurityAuthenticator::Open => fidl_security::Authentication {
                protocol: fidl_security::Protocol::Open,
                credentials: None,
            },
            SecurityAuthenticator::Wep(wep) => fidl_security::Authentication {
                protocol: fidl_security::Protocol::Wep,
                credentials: Some(Box::new(fidl_security::Credentials::Wep(wep.into()))),
            },
            SecurityAuthenticator::Wpa(wpa) => {
                use wpa::Authentication::{Enterprise, Personal};
                use wpa::Wpa::{Wpa1, Wpa2, Wpa3};

                let protocol = match (&wpa, wpa.to_credentials()) {
                    (Wpa1 { .. }, _) => fidl_security::Protocol::Wpa1,
                    (Wpa2 { .. }, Personal(_)) => fidl_security::Protocol::Wpa2Personal,
                    (Wpa2 { .. }, Enterprise(_)) => fidl_security::Protocol::Wpa2Enterprise,
                    (Wpa3 { .. }, Personal(_)) => fidl_security::Protocol::Wpa3Personal,
                    (Wpa3 { .. }, Enterprise(_)) => fidl_security::Protocol::Wpa3Enterprise,
                };
                fidl_security::Authentication {
                    protocol,
                    // TODO(https://fxbug.dev/42174395): This panics when encountering WPA Enterprise.
                    credentials: Some(Box::new(fidl_security::Credentials::Wpa(
                        wpa.into_credentials().into(),
                    ))),
                }
            }
        }
    }
}

/// Converts an `Authentication` FIDL datagram into a `SecurityAuthenticator`.
///
/// This conversion should be preferred where possible.
///
/// # Errors
///
/// Returns an error if the `Authentication` datagram is invalid, such as specifying contradictory
/// protocols or encoding incompatible or invalid credentials.
impl TryFrom<fidl_security::Authentication> for SecurityAuthenticator {
    type Error = SecurityError;

    fn try_from(authentication: fidl_security::Authentication) -> Result<Self, Self::Error> {
        let fidl_security::Authentication { protocol, credentials } = authentication;
        match protocol {
            fidl_security::Protocol::Open => match credentials {
                None => Ok(SecurityAuthenticator::Open),
                _ => Err(SecurityError::Incompatible),
            },
            fidl_security::Protocol::Wep => credentials
                .ok_or_else(|| SecurityError::Incompatible)? // No credentials.
                .into_wep()
                .map(wep::WepAuthenticator::try_from)
                .transpose()? // Conversion failure.
                .map(From::from)
                .ok_or_else(|| SecurityError::Incompatible), // Non-WEP credentials.
            fidl_security::Protocol::Wpa1 => credentials
                .ok_or_else(|| SecurityError::Incompatible)? // No credentials.
                .into_wpa()
                .map(wpa::Wpa1Credentials::try_from)
                .transpose()? // Conversion failure.
                .map(|credentials| wpa::WpaAuthenticator::Wpa1 { credentials })
                .map(From::from)
                .ok_or_else(|| SecurityError::Incompatible), // Non-WPA credentials.
            fidl_security::Protocol::Wpa2Personal => credentials
                .ok_or_else(|| SecurityError::Incompatible)? // No credentials.
                .into_wpa()
                .map(wpa::Wpa2PersonalCredentials::try_from)
                .transpose()? // Conversion failure.
                .map(From::from)
                .map(|authentication| wpa::WpaAuthenticator::Wpa2 { cipher: None, authentication })
                .map(From::from)
                .ok_or_else(|| SecurityError::Incompatible), // Non-WPA credentials.
            fidl_security::Protocol::Wpa3Personal => credentials
                .ok_or_else(|| SecurityError::Incompatible)? // No credentials.
                .into_wpa()
                .map(wpa::Wpa3PersonalCredentials::try_from)
                .transpose()? // Conversion failure.
                .map(From::from)
                .map(|authentication| wpa::WpaAuthenticator::Wpa3 { cipher: None, authentication })
                .map(From::from)
                .ok_or_else(|| SecurityError::Incompatible), // Non-WPA credentials.
            // TODO(https://fxbug.dev/42174395): This returns an error when encountering WPA Enterprise
            //                        protocols. Some conversions of composing types panic, but
            //                        this top-level conversion insulates client code from this and
            //                        instead yields an error.
            _ => Err(SecurityError::Incompatible),
        }
    }
}

#[cfg(test)]
mod tests {
    use fidl_fuchsia_wlan_common_security as fidl_security;

    use test_case::test_case;

    use crate::security::{
        wpa::{self, Authentication, Wpa2PersonalCredentials},
        SecurityAuthenticator, SecurityError,
    };

    pub trait AuthenticationTestCase: Sized {
        fn wpa2_personal_psk() -> Self;
        fn wpa3_personal_psk() -> Self;
        fn wpa3_personal_wep_key() -> Self;
        fn wpa3_personal_no_credentials() -> Self;
    }

    impl AuthenticationTestCase for fidl_security::Authentication {
        fn wpa2_personal_psk() -> Self {
            fidl_security::Authentication {
                protocol: fidl_security::Protocol::Wpa2Personal,
                credentials: Some(Box::new(fidl_security::Credentials::Wpa(
                    fidl_security::WpaCredentials::Psk([0u8; 32]),
                ))),
            }
        }

        // Invalid: WPA3 with PSK.
        fn wpa3_personal_psk() -> Self {
            fidl_security::Authentication {
                protocol: fidl_security::Protocol::Wpa3Personal,
                credentials: Some(Box::new(fidl_security::Credentials::Wpa(
                    fidl_security::WpaCredentials::Psk([0u8; 32]),
                ))),
            }
        }

        // Invalid: WPA3 with WEP key.
        fn wpa3_personal_wep_key() -> Self {
            fidl_security::Authentication {
                protocol: fidl_security::Protocol::Wpa3Personal,
                credentials: Some(Box::new(fidl_security::Credentials::Wep(
                    fidl_security::WepCredentials { key: vec![0u8; 13] },
                ))),
            }
        }

        // Invalid: WPA3 with no credentials.
        fn wpa3_personal_no_credentials() -> Self {
            fidl_security::Authentication {
                protocol: fidl_security::Protocol::Wpa3Personal,
                credentials: None,
            }
        }
    }

    // TODO(seanolson): Move this assertion into a `SecurityAuthenticatorAssertion` trait (a la
    //                  `AuthenticationTestCase`) and test via the `using` pattern in the
    //                  `test-case` 2.0.0 series.
    #[test_case(AuthenticationTestCase::wpa2_personal_psk() => matches
        Ok(SecurityAuthenticator::Wpa(wpa::Wpa::Wpa2 {
            authentication: Authentication::Personal(Wpa2PersonalCredentials::Psk(_)),
            ..
        }))
    )]
    #[test_case(AuthenticationTestCase::wpa3_personal_psk() => Err(SecurityError::Incompatible))]
    #[test_case(AuthenticationTestCase::wpa3_personal_wep_key() => Err(SecurityError::Incompatible))]
    #[test_case(AuthenticationTestCase::wpa3_personal_no_credentials() => Err(SecurityError::Incompatible))]
    fn security_authenticator_from_authentication_fidl(
        authentication: fidl_security::Authentication,
    ) -> Result<SecurityAuthenticator, SecurityError> {
        SecurityAuthenticator::try_from(authentication)
    }

    #[test]
    fn authentication_fidl_from_security_authenticator() {
        let authenticator = SecurityAuthenticator::Wpa(wpa::WpaAuthenticator::Wpa3 {
            authentication: wpa::Authentication::Personal(
                wpa::Wpa3PersonalCredentials::Passphrase("roflcopter".try_into().unwrap()),
            ),
            cipher: None,
        });
        let authentication = fidl_security::Authentication::from(authenticator);
        assert_eq!(
            authentication,
            fidl_security::Authentication {
                protocol: fidl_security::Protocol::Wpa3Personal,
                credentials: Some(Box::new(fidl_security::Credentials::Wpa(
                    fidl_security::WpaCredentials::Passphrase(b"roflcopter".as_slice().into()),
                ))),
            }
        );
    }
}