use super::*;
use rsa::{PublicKeyParts, RSAPublicKey};
use std::convert::TryInto;

impl KeyExt for RSAPublicKey {
    fn to_pem_pkcs8(&self) -> Result<String, KeyError> {
        let bytes = write_pkcs1(self);
        let oid = yasna::models::ObjectIdentifier::from_slice(&RSA_OID);

        let contents = yasna::construct_der(|writer| {
            writer.write_sequence(|writer| {
                writer.next().write_sequence(|writer| {
                    writer.next().write_oid(&oid);
                    writer.next().write_null();
                });
                writer
                    .next()
                    .write_bitvec(&bit_vec::BitVec::from_bytes(&bytes));
            });
        });

        let p = pem::Pem {
            tag: "PUBLIC KEY".to_owned(),
            contents,
        };

        Ok(pem::encode(&p))
    }

    fn from_pem_pkcs8(pem: &str) -> Result<Self, KeyError> {
        let data = pem::parse(pem).map_err(|_| KeyError::Pem)?;

        if data.tag != "PUBLIC KEY" {
            return Err(KeyError::Kind);
        }

        data.try_into().map_err(KeyError::Parse)
    }

    fn to_pem_pkcs1(&self) -> Result<String, KeyError> {
        let contents = write_pkcs1(self);

        let p = pem::Pem {
            tag: "RSA PUBLIC KEY".to_owned(),
            contents,
        };

        Ok(pem::encode(&p))
    }

    fn from_pem_pkcs1(pem: &str) -> Result<Self, KeyError> {
        let data = pem::parse(pem).map_err(|_| KeyError::Pem)?;

        if data.tag != "RSA PUBLIC KEY" {
            return Err(KeyError::Kind);
        }

        data.try_into().map_err(KeyError::Parse)
    }
}

fn write_pkcs1(rsa: &RSAPublicKey) -> Vec<u8> {
    yasna::construct_der(|writer| {
        writer.write_sequence(|writer| {
            writer.next().write_biguint(&from_dig(rsa.n()));
            writer.next().write_biguint(&from_dig(rsa.e()));
        })
    })
}