bitbottle 0.10.0

use std::io::Read;
use std::path::PathBuf;
use crate::asymmetric::BottleSecretKey;
use crate::bottle::{BottleReader, BottleStream};
use crate::bottle_cap::{BottleCap, BottleType};
use crate::bottle_error::{BottleError, BottleResult};
use crate::compressed_bottle::CompressedBottleReader;
use crate::compression::CompressionAlgorithm;
use crate::encrypted_bottle::{
    EncryptedBottleInfo, EncryptedBottleReader, EncryptedBottleReaderOptions,
    EncryptionKey
};
use crate::file_atlas::FileAtlasRef;
use crate::file_bottle::FileListBottleReader;
use crate::file_list::{FileList, FileListRef, Symlink};
use crate::hashing::HashType;
use crate::signed_bottle::{SignatureAlgorithm, SignedBottleReader};
use crate::streams::{drain_stream, ReadStreamRef, StreamBottle, WrappedRead};
use crate::BottlePublicKey;


pub enum ArchiveReaderEvent {
    // error reading archive
    Error(BottleError),

    // metadata (headers) from the next bottle -- always matched with a later BottleEnd
    BottleCap(BottleCap),

    // signed bottle
    Signed {
        algorithm: SignatureAlgorithm,
        public_key: Vec<u8>,
        signed_by: String,
        verify_key_name: Option<String>,
    },

    // encrypted bottle
    Encrypted(EncryptedBottleInfo),

    // compressed bottle
    Compressed(CompressionAlgorithm),

    // file list events: first a bunch of FileAtlas, then FileListDone,
    // then a bunch of FileBlock, then BottleEnd
    FileListStart { hash_type: HashType, block_count: usize, file_count: usize },
    FileAtlas(FileAtlasRef),
    FileListDone { file_list: FileListRef, bad_path_list: Vec<(PathBuf, PathBuf)>, stray_symlinks: Vec<Symlink> },
    FileBlock { size: usize, hash: Vec<u8>, data: Option<Vec<u8>>, bottle_cap: BottleCap },

    // when extracting, a FileBlock event becomes FileBlockWritten
    FileBlockWritten { size: usize, write_count: usize },

    // found a bottle i don't understand
    UnknownBottleType,
    UnknownBottleData { size: usize, data: Option<Vec<u8>> },

    // end of that bottle
    BottleEnd(BottleType),

    // drained all bottles successfully
    Done,
}


pub struct ArchiveReaderOptions {
    // read each file block into memory (for unpacking or verifying), or just
    // report the size and hash?
    pub read_blocks: bool,

    // process known bottle types (compressed, encrypted, and so on), or just
    // dump each bottle type as "unknown" as if we have no idea what it is?
    pub process_bottles: bool,

    // if it's signed, you can supply public keys that are allowed:
    pub verify_keys: Vec<Box<dyn BottlePublicKey>>,

    // if it's encrypted, you can supply secret keys to try, or a password:
    pub secret_keys: Vec<Box<dyn BottleSecretKey>>,
    pub encryption_key: Option<EncryptionKey>,

    // when draining an unknown bottle type (very unlikely), how many bytes
    // are you willing to receive from unknown data bottles?
    pub max_unknown_data_capture: usize,

    /// Allow older encrypted archives that are missing a key commitment in
    /// the header?
    pub allow_missing_key_commitment: bool,
}

impl ArchiveReaderOptions {
    pub fn new() -> ArchiveReaderOptions {
        ArchiveReaderOptions {
            read_blocks: true,
            process_bottles: true,
            verify_keys: vec![],
            secret_keys: vec![],
            encryption_key: None,
            max_unknown_data_capture: 128,
            allow_missing_key_commitment: false,
        }
    }
}

impl Default for ArchiveReaderOptions {
    fn default() -> ArchiveReaderOptions {
        ArchiveReaderOptions::new()
    }
}


enum State {
    // emit the bottle cap first:
    NewBottle,
    // reading streams:
    Streams,
    // unknown bottle type, so we're just draining it:
    Draining,
    // reached the end of an opaque bottle
    EndBottle,
}


pub struct ArchiveReader {
    options: ArchiveReaderOptions,

    // stack of the `io::Read` of each bottle as we open them
    stack: Vec<(ReadStreamRef, BottleType)>,

    // when streaming out of the top stack bottle:
    bottle_type: BottleType,
    current_bottle: Option<BottleReader<ReadStreamRef>>,

    // when processing the file list at the core:
    file_bottle: Option<FileListBottleReader<ReadStreamRef>>,
    // collected files and block index:
    file_list: FileListRef,
    // list of files with bad paths:
    bad_path_list: Vec<(PathBuf, PathBuf)>,
    // notice the switch from files to blocks:
    in_blocks: bool,
    // if there were no acceptable signatures, bail on the next iteration
    fail_bad_signature: bool,

    // stack of unknown bottles we're draining when we're in unknown types
    reader_stack: Vec<BottleReader<ReadStreamRef>>,

    state: State,
}

impl ArchiveReader {
    pub fn new(
        reader: Box<dyn Read>,
        options: ArchiveReaderOptions,
    ) -> BottleResult<ArchiveReader> {
        // wrap the "Read" into a "StreamingBottle" where finish() does nothing:
        let stream = WrappedRead::new(reader);
        Ok(ArchiveReader {
            options,
            stack: vec![ (stream.stream(), BottleType::Test) ],
            bottle_type: BottleType::Test,
            current_bottle: None,
            file_bottle: None,
            file_list: FileList::new().to_shared(),
            bad_path_list: Vec::new(),
            in_blocks: false,
            fail_bad_signature: false,
            reader_stack: Vec::new(),
            state: State::NewBottle,
        })
    }

    fn push_bottle(&mut self, next_bottle: ReadStreamRef, bottle_type: BottleType) {
        self.stack.push((next_bottle, bottle_type));
        self.state = State::NewBottle;
    }

    fn pop_bottle(&mut self) -> BottleResult<BottleType> {
        let (mut stream, bottle_type) = self.stack.pop().unwrap();
        stream.finish()?;
        Ok(bottle_type)
    }

    fn process_file_list(
        &mut self,
        mut file_bottle: FileListBottleReader<ReadStreamRef>
    ) -> BottleResult<ArchiveReaderEvent> {
        if !file_bottle.is_in_files() {
            if !self.in_blocks {
                self.in_blocks = true;
                self.file_list.borrow_mut().build_file_map();
                let mut bad_path_list: Vec<(PathBuf, PathBuf)> = Vec::new();
                bad_path_list.append(&mut self.bad_path_list);
                let stray_symlinks = self.file_list.borrow_mut().drop_stray_symlinks();
                self.file_bottle = Some(file_bottle);
                Ok(ArchiveReaderEvent::FileListDone {
                    file_list: self.file_list.clone(),
                    bad_path_list,
                    stray_symlinks,
                })
            } else if let Some(block) = file_bottle.next_block()? {
                self.file_bottle = Some(file_bottle);
                Ok(ArchiveReaderEvent::FileBlock {
                    size: block.block.size,
                    hash: block.block.hash.to_vec(),
                    data: self.options.read_blocks.then_some(block.data),
                    bottle_cap: block.bottle_cap,
                })
            } else {
                file_bottle.close()?;
                self.state = State::EndBottle;
                Ok(ArchiveReaderEvent::BottleEnd(BottleType::FileList))
            }
        } else if let Some(atlas) = file_bottle.next_file()? {
            let bad_path = atlas.borrow_mut().fix_bad_path();
            if let Some(bad_path) = bad_path {
                self.bad_path_list.push((bad_path, atlas.borrow().normalized_path.clone()));
            }
            self.file_list.borrow_mut().files.push(atlas.clone());
            self.file_bottle = Some(file_bottle);
            Ok(ArchiveReaderEvent::FileAtlas(atlas.clone()))
        } else {
            Err(BottleError::InvalidBottleState)
        }
    }

    fn open_bottle(&mut self, mut bottle_reader: BottleReader<ReadStreamRef>) -> BottleResult<ArchiveReaderEvent> {
        if !self.options.process_bottles {
            self.state = State::Draining;
            self.current_bottle = Some(bottle_reader);
            return Ok(ArchiveReaderEvent::UnknownBottleType);
        }

        let event = match self.bottle_type {
            BottleType::Signed => {
                let reader = SignedBottleReader::new(bottle_reader)?;
                let algorithm = reader.algorithm;
                let public_key = reader.public_key.as_bytes(false).to_vec();
                let signed_by = reader.public_key.name().to_string();
                let verify_keys = &self.options.verify_keys;
                if !verify_keys.is_empty() && !verify_keys.iter().any(|vk| vk.as_bytes(false) == public_key) {
                    self.fail_bad_signature = true;
                }
                let verify_key_name = verify_keys.iter()
                    .find(|vk| vk.as_bytes(false) == public_key)
                    .map(|pk| pk.name().to_string());
                self.push_bottle(reader.stream(), BottleType::Signed);
                ArchiveReaderEvent::Signed { algorithm, public_key, signed_by, verify_key_name }
            }

            BottleType::Encrypted => {
                let info = EncryptedBottleReader::unpack_info(&mut bottle_reader)?;

                // decrypt if possible
                let encryption_key = self.options.encryption_key.clone().unwrap_or(EncryptionKey::Generated);
                let reader = EncryptedBottleReader::build_with_info(
                    bottle_reader,
                    &info,
                    EncryptedBottleReaderOptions {
                        key_type: encryption_key,
                        possible_keys: Some(&self.options.secret_keys),
                        allow_missing_key_commitment: self.options.allow_missing_key_commitment,
                    },
                )?;
                self.push_bottle(reader.stream(), BottleType::Encrypted);

                ArchiveReaderEvent::Encrypted(info)
            },

            BottleType::Compressed => {
                let reader = CompressedBottleReader::new(bottle_reader)?;
                let algorithm = reader.algorithm;
                self.push_bottle(reader.stream(), BottleType::Compressed);
                ArchiveReaderEvent::Compressed(algorithm)
            },

            // but this is why we're all here:
            BottleType::FileList => {
                let file_bottle = FileListBottleReader::new(bottle_reader, self.options.read_blocks)?;
                let event = ArchiveReaderEvent::FileListStart {
                    hash_type: file_bottle.hash_type,
                    block_count: file_bottle.block_count,
                    file_count: file_bottle.file_count,
                };
                self.file_bottle = Some(file_bottle);
                self.file_list.borrow_mut().clear();
                self.in_blocks = false;
                event
            },

            _ => {
                self.state = State::Draining;
                self.current_bottle = Some(bottle_reader);
                ArchiveReaderEvent::UnknownBottleType
            },
        };
        Ok(event)
    }

    fn drain_bottle(&mut self, mut bottle_reader: BottleReader<ReadStreamRef>) -> BottleResult<ArchiveReaderEvent> {
        match bottle_reader.next_stream()? {
            BottleStream::Data => {
                let (size, data) = drain_stream(bottle_reader.data_stream()?, self.options.max_unknown_data_capture)?;
                bottle_reader.close_stream()?;
                self.current_bottle = Some(bottle_reader);
                Ok(ArchiveReaderEvent::UnknownBottleData { size, data })
            },
            BottleStream::Bottle => {
                let reader = bottle_reader.take_bottle_reader()?;
                let bottle_cap = reader.bottle_cap.clone();
                self.current_bottle = Some(*reader);
                self.reader_stack.push(bottle_reader);
                Ok(ArchiveReaderEvent::BottleCap(bottle_cap))
            },
            BottleStream::End => {
                let bottle_type = bottle_reader.bottle_cap.bottle_type;
                bottle_reader.close()?;
                if !self.reader_stack.is_empty() {
                    self.current_bottle = Some(self.reader_stack.pop().unwrap());
                } else {
                    self.state = State::EndBottle;
                }
                Ok(ArchiveReaderEvent::BottleEnd(bottle_type))
            },
        }
    }
}

impl Iterator for ArchiveReader {
    type Item = ArchiveReaderEvent;

    fn next(&mut self) -> Option<ArchiveReaderEvent> {
        if self.stack.is_empty() { return None }
        if self.fail_bad_signature { return Some(ArchiveReaderEvent::Error(BottleError::NoAcceptableSignature)); }

        let event = if let Some(file_bottle) = self.file_bottle.take() {
            self.process_file_list(file_bottle)
        } else {
            match self.state {
                State::NewBottle => {
                    self.state = State::Streams;
                    BottleReader::new(self.stack.last_mut().unwrap().0.clone()).map(|bottle_reader| {
                        let bottle_cap = bottle_reader.bottle_cap.clone();
                        self.bottle_type = bottle_cap.bottle_type;
                        self.current_bottle = Some(bottle_reader);
                        ArchiveReaderEvent::BottleCap(bottle_cap)
                    })
                },

                State::Streams => {
                    let bottle_reader = self.current_bottle.take().unwrap();
                    self.open_bottle(bottle_reader)
                },

                State::Draining => {
                    let bottle_reader = self.current_bottle.take().unwrap();
                    self.drain_bottle(bottle_reader)
                },

                State::EndBottle => {
                    self.pop_bottle().map(|bottle_type| {
                        if self.stack.is_empty() {
                            ArchiveReaderEvent::Done
                        } else {
                            // each will correspond to one BottleCap
                            ArchiveReaderEvent::BottleEnd(bottle_type)
                        }
                    })
                },
            }
        };

        Some(event.unwrap_or_else(|e| {
            self.current_bottle = None;
            self.stack.clear();
            ArchiveReaderEvent::Error(e)
        }))
    }
}