// Panopticon - A libre program analysis library for machine code
// Copyright (C) 2014-2018  The Panopticon Developers
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

use std::ops::Range;
use std::io::Cursor;
use std::fs::File;
use std::sync::Arc;
use std::path::{Path, PathBuf};

use byteorder::{
    ReadBytesExt,
    BigEndian,
    LittleEndian
};
use memmap::Mmap;

use {
    Result,
    Str,
    Endianess,
    Uuid,
};

#[derive(Clone, Debug)]
/// A continuous address space. Regions are an array of cells numbered from 0. Each cell is either
/// `undefined` of carries a single byte value.
pub enum Region {
    /// All undefined values
    Undefined{
        /// Addressable space. All cells in a region have addresses between 0 and `2^address_bits -
        /// 1`.
        address_bits: usize,
        /// Human readable name. This name can vary over the livetime of a region and bears no
        /// semantic meaning. It's for use in UIs.
        name: Str,
        /// Immutable UUID. The UUID of a region never changes. Can be used for housekeeping and
        /// associating additional data with this region.
        uuid: Uuid,
    },
    /// In memory buffer
    Buffer{
        /// Addressable space. All cells in a region have addresses between 0 and `2^address_bits -
        /// 1`.
        address_bits: usize,
        /// Human readable name. This name can vary over the livetime of a region and bears no
        /// semantic meaning. It's for use in UIs.
        name: Str,
        /// Immutable UUID. The UUID of a region never changes. Can be used for housekeeping and
        /// associating additional data with this region.
        uuid: Uuid,
        /// Positon of the buffer in the address space.
        offset: u64,
        /// Region contents
        buffer: Arc<[u8]>,
    },
    /// Memory mapped file.
    File{
        /// Addressable space. All cells in a region have addresses between 0 and `2^address_bits -
        /// 1`.
        address_bits: usize,
        /// Human readable name. This name can vary over the livetime of a region and bears no
        /// semantic meaning. It's for use in UIs.
        name: Str,
        /// Immutable UUID. The UUID of a region never changes. Can be used for housekeeping and
        /// associating additional data with this region.
        uuid: Uuid,
        /// Positon of the file in the address space.
        offset: u64,
        /// mmap()'d file.
        file: Arc<Mmap>,
        /// Path to the mapped file
        path: Option<PathBuf>,
        /// Start of mapping inside file
        file_offset: u64,
    },
}

impl Region {
    /// Creates a new completly undefined region with `name` and `address_bits`.
    pub fn undefined<S: Into<Str>>(name: S, address_bits: usize) -> Region {
        Region::Undefined{
            address_bits: address_bits,
            name: name.into(),
            uuid: Uuid::now(),
        }
    }

    /// Creates a new region called `name` that is `address_bits` large. Maps `mmap` to `offset`.
    pub fn from_mmap<S, O, P>(name: S, address_bits: usize, mmap: Mmap, path: P, file_offset: O, offset: O) -> Region where S: Into<Str>, O: Into<Option<u64>>, P: Into<Option<PathBuf>> {
        Region::File{
            offset: offset.into().unwrap_or(0),
            name: name.into(),
            file: Arc::new(mmap),
            address_bits: address_bits,
            uuid: Uuid::now(),
            path: path.into(),
            file_offset: file_offset.into().unwrap_or(0),
        }
    }

    /// Creates a new region called `name` that is `address_bits` large. Maps `fd` to `offset`.
    pub fn from_file<S, O, P>(name: S, address_bits: usize, fd: File, path: P, offset: O) -> Result<Region> where S: Into<Str>, O: Into<Option<u64>>, P: Into<Option<PathBuf>> {
        Ok(Region::File{
            offset: offset.into().unwrap_or(0),
            name: name.into(),
            file: Arc::new(unsafe { Mmap::map(&fd)? }),
            address_bits: address_bits,
            uuid: Uuid::now(),
            path: path.into(),
            file_offset: 0,
        })
    }

    /// Creates a new region called `name` that is `address_bits` large. Maps `buf` to `offset`.
    pub fn from_buf<S: Into<Str>, O: Into<Option<u64>>,B: Into<Arc<[u8]>>>(name: S, address_bits: usize, buf: B, offset: O) -> Region {
        Region::Buffer{
            offset: offset.into().unwrap_or(0),
            name: name.into(),
            buffer: buf.into(),
            address_bits: address_bits,
            uuid: Uuid::now(),
        }
    }

    /// If this is a mmap()'d file, return the path to it and mmap() starting position.
    pub fn file<'a>(&'a self) -> Option<(&'a Path, u64)> {
        match self {
            &Region::Undefined{ .. } => None,
            &Region::File{ ref path, file_offset, .. } =>
                path.as_ref().map(|x| (PathBuf::as_path(x), file_offset)),
            &Region::Buffer{ .. } => None,
        }
    }

    /// Returns the human readable name of this region.
    pub fn name<'a>(&'a self) -> &'a Str {
        match self {
            &Region::Undefined{ ref name,.. } => name,
            &Region::File{ ref name,.. } => name,
            &Region::Buffer{ ref name,.. } => name,
        }
    }

    /// Changes the human readable name to `new`.
    pub fn rename(&mut self, new: Str) {
        match self {
            &mut Region::Undefined{ ref mut name,.. } => { *name = new }
            &mut Region::File{ ref mut name,.. } => { *name = new }
            &mut Region::Buffer{ ref mut  name,.. } => { *name = new }
        }
    }

    /// Returns the immutable UUID of this region.
    pub fn uuid<'a>(&'a self) -> &'a Uuid {
        match self {
            &Region::Undefined{ ref uuid,.. } => uuid,
            &Region::File{ ref uuid,.. } => uuid,
            &Region::Buffer{ ref uuid,.. } => uuid,
        }
    }

    /// Size of this region.
    pub fn address_bits(&self) -> usize {
        match self {
            &Region::Undefined{ address_bits,.. } => address_bits,
            &Region::File{ address_bits,.. } => address_bits,
            &Region::Buffer{ address_bits,.. } => address_bits,
        }
    }

    /// Defined range
    pub fn defined(&self) -> Range<u64> {
        match self {
            &Region::Undefined{ .. } => 0..0,
            &Region::File{ offset, ref file,.. } =>
                offset..offset + file.len() as u64,
            &Region::Buffer{ offset, ref buffer,.. } =>
                offset..offset + buffer.len() as u64,
        }
    }

    /// Fill `buf` the the values starting at `address`. Fails if `address..address + buf.len()`
    /// is outside of the addressable range or contains undefined values.
    pub fn read<'a>(&'a self, start: u64, len: usize) -> Result<&'a[u8]> {
        if !self.in_range(start..start + len as u64) { return Err("Out of range".into()); }
        if len == 0 || !self.is_defined(start..start + len as u64) {
            return Err("Undefined".into());
        }

        match self {
            &Region::Undefined{ .. } => unreachable!(),
            &Region::File{ ref file,.. } => {
                let a = (start as usize)..(start as usize + len);
                Ok(&file[a])
            }
            &Region::Buffer{ ref buffer,.. } => {
                let a = (start as usize)..(start as usize + len);
                Ok(&buffer[a])
            }
        }
    }

    /// Trys to fill `buf` the the values starting at `address`. Returns early if `address..address + buf.len()`
    /// contains undefined values and fails if it's outside of the addressable range. Returns the number of bytes read.
    pub fn try_read(&self, address: u64, buf: &mut [u8]) -> Result<usize> {
        use std::cmp;

        if !self.in_range(address..address + 1) { return Err("Out of range".into()); }
        if !self.is_defined(address..address + 1) { return Ok(0); }

        match self {
            &Region::Undefined{ .. } if buf.len() == 0 => Ok(0),
            &Region::Undefined{ .. } => unreachable!(),
            &Region::File{ ref file, offset,.. } => {
                let o = offset as usize;
                let l = cmp::min(buf.len(), file.len());
                let a = (address as usize - o)..(address as usize + l - o);

                buf[0..l].clone_from_slice(&file[a]);
                Ok(l)
            }

            &Region::Buffer{ ref buffer, offset,.. } => {
                let o = offset as usize;
                let address = address as usize - o;
                let l = cmp::min(buf.len(), buffer.len() - address);
                let a = (address as usize)..(address as usize + l);

                buf[0..l].clone_from_slice(&buffer[a]);
                Ok(l)
            }
        }
    }

    /// Reads an `bytes` large integer from `address` and zero-extends it to an `u64`. Fails if `address..address + bytes`
    /// is outside of the addressable range or contains undefined values.
    pub fn read_integer(&self, address: u64, endianess: Endianess, bytes: usize) -> Result<u64> {
        match bytes {
            1 => {
                Ok(self.read(address, 1)?[0] as u64)
            }
            2 => {
                let buf = self.read(address, 2)?;
                let mut cur = Cursor::new(buf);

                match endianess {
                    Endianess::Little => Ok(cur.read_u16::<LittleEndian>()? as u64),
                    Endianess::Big => Ok(cur.read_u16::<BigEndian>()? as u64),
                }
            }
            4 => {
                let buf = self.read(address, 4)?;
                let mut cur = Cursor::new(buf);

                match endianess {
                    Endianess::Little => Ok(cur.read_u32::<LittleEndian>()? as u64),
                    Endianess::Big => Ok(cur.read_u32::<BigEndian>()? as u64),
                }
            }
            8 => {
                let buf = self.read(address, 8)?;
                let mut cur = Cursor::new(buf);

                match endianess {
                    Endianess::Little => cur.read_u64::<LittleEndian>().map_err(|e| e.into()),
                    Endianess::Big => cur.read_u64::<BigEndian>().map_err(|e| e.into()),
                }
            }
            _ => Err(format!("reaidng a {} byte integer is unimplemented",bytes).into()),
        }
    }

    /// Returns true if range is in the addressable space of this region, i.e. `0..address_bits^2 -
    /// `.
    pub fn in_range(&self, range: Range<u64>) -> bool {
        match self {
            &Region::Undefined{ address_bits,.. } => {
                64 - range.start.leading_zeros() as usize <= address_bits &&
                    64 - range.end.saturating_sub(1).leading_zeros() as usize <= address_bits
            }
            &Region::Buffer{ address_bits,.. } => {
                64 - range.start.leading_zeros() as usize <= address_bits &&
                    64 - range.end.saturating_sub(1).leading_zeros() as usize <= address_bits
            }
            &Region::File{ address_bits,.. } => {
                64 - range.start.leading_zeros() as usize <= address_bits &&
                    64 - range.end.saturating_sub(1).leading_zeros() as usize <= address_bits
            }
        }
    }

    /// Returns true if range is in the addressable space of this region and contains no undefined
    /// values.
    pub fn is_defined(&self, range: Range<u64>) -> bool {
        match self {
            &Region::Undefined{ .. } => false,
            &Region::Buffer{ ref buffer, offset, address_bits,.. } => {
                let end = if address_bits < 64 { 1 << address_bits } else { 0xFFFFFFFFFFFFFFFF };
                let head = 0..offset;
                let tail = offset + (buffer.len() as u64)..end;
                let outside = head.end > range.start || tail.start < range.end;

                !outside
            }
            &Region::File{ ref file, offset, address_bits,.. } => {
                let end = if address_bits < 64 { 1 << address_bits } else { 0xFFFFFFFFFFFFFFFF };
                let head = 0..offset;
                let tail = offset + (file.len() as u64)..end;
                let outside = head.end > range.start || tail.start < range.end;

                !outside
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn construct() {
        let b: Box<[u8]> = vec![1, 2, 3].into();
        let l1 = Region::undefined("l1",6);
        let l2 = Region::from_buf("l2",3,b,None);

        assert_eq!(l1.defined(), 0..0);
        assert_eq!(l2.defined(), 0..3);
        assert_eq!(l1.address_bits(), 6);
        assert_eq!(l1.name(), "l1");
        assert_eq!(l2.address_bits(), 3);
        assert_eq!(l2.name(), "l2");
    }

    #[test]
    fn is_defined() {
        let b: Box<[u8]> = vec![1, 2, 3].into();
        let l1 = Region::from_buf("l2",16,b,5);

        assert_eq!(l1.defined(), 5..8);
        assert_eq!(l1.is_defined(0..1), false);
        assert!(l1.try_read(0, &mut vec![0u8; 1]).is_ok());
        assert_eq!(l1.is_defined(0..5), false);
        assert!(l1.try_read(0, &mut vec![0u8; 4]).is_ok());
        assert_eq!(l1.is_defined(3..7), false);
        assert!(l1.try_read(3, &mut vec![0u8; 4]).is_ok());
        assert_eq!(l1.is_defined(5..6), true);
        assert!(l1.try_read(5, &mut vec![0u8; 1]).is_ok());
        assert_eq!(l1.is_defined(5..8), true);
        assert!(l1.try_read(5, &mut vec![0u8; 3]).is_ok());
        assert_eq!(l1.is_defined(9..20), false);
        assert!(l1.try_read(9, &mut vec![0u8; 10]).is_ok());
    }

    #[test]
    fn in_range() {
        let b: Box<[u8]> = vec![1, 2, 3].into();
        let l1 = Region::from_buf("l2",16,b,0);

        assert_eq!(l1.defined(), 0..3);
        assert_eq!(l1.in_range(0..1), true);
        assert!(l1.try_read(0, &mut vec![0u8; 1]).is_ok());
        assert_eq!(l1.in_range(0x10000..0x10001), false);
        assert_eq!(l1.in_range(3..0x1_00_00), true);
        assert_eq!(l1.in_range(5..0x10001), false);
    }

    #[test]
    fn read() {
        let b: Box<[u8]> = vec![1, 2, 3].into();
        let l1 = Region::from_buf("l2",3,b,None);

        {
            let buf = l1.read(0, 1).unwrap();
            assert_eq!(buf, [1]);
        }

        {
            let buf = l1.read(2, 1).unwrap();
            assert_eq!(buf, [3]);
        }

        {
            let buf = l1.read(1, 2).unwrap();
            assert_eq!(buf, [2,3]);
        }
    }

    #[test]
    fn try_read() {
        let b: Box<[u8]> = vec![1, 2, 3].into();
        let l1 = Region::from_buf("l2",3,b,None);

        {
            let mut buf = [0u8; 3];
            assert_eq!(l1.try_read(0, &mut buf).ok(), Some(3));
            assert_eq!(buf, [1, 2, 3]);
        }

        {
            let mut buf = [0u8; 3];
            assert_eq!(l1.try_read(2, &mut buf).ok(), Some(1));
            assert_eq!(buf, [3, 0, 0]);
        }

        {
            let mut buf = [0u8; 3];
            assert_eq!(l1.try_read(1, &mut buf).ok(), Some(2));
            assert_eq!(buf, [2, 3, 0]);
        }
    }
}
