// 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::borrow::Cow;
use smallvec::SmallVec;
use {
    Region,
    Result,
    Value,
    Variable,
    Statement,
    Operation,
    FlowOperation,
    MemoryOperation,
    Guard,
    Str,
    Names,
    Segments,
    Segment,
    Name,
    Endianess,
    Strings,
};

/// CPU architecture and instruction set.
pub trait Architecture: Clone {
    /// This type can describes the CPU state. For x86 this would be the mode, for ARM whenever
    /// Thumb is active.
    type Configuration: Clone + Send;

    /// Given a memory image and a configuration the functions extracts a set of entry points.
    /// # Return
    /// Tuples of entry point name, offset form the start of the region and optional comment.
    fn prepare(&Region, &Self::Configuration) -> Result<Cow<'static, [(&'static str, u64, &'static str)]>>;

    /// Start to disassemble a single Opcode inside a given region at a given address.
    fn decode(&Region, u64, &Self::Configuration, names: &mut Names,
              segments: &mut Segments, strings: &mut Strings,
              matches: &mut Vec<Match>) -> Result<()>;
}

/// Result of a single disassembly operation.
#[derive(Debug,Clone)]
pub struct Match {
    /// Recognized address range
    pub area: Range<u64>,
    /// Matched mnemonic.
    pub opcode: Str,
    /// Opcode arguments.
    pub operands: SmallVec<[Value; 3]>,
    //pub format_string: SmallVec<[MnemonicFormatToken; 5]>,
    /// IL statements modeling the semantics.
    pub instructions: Vec<Statement>,

    /// Jumps/branches originating from the recovered mnemonics
    pub jumps: SmallVec<[(u64, Value, Guard); 2]>,
}

/// ISA used for testing.
/// Single digits are interpreted as 32 bit constants, lower case letters as 32 bit registers or 1
/// bit flags depending on the context. First argument is the result. Upper case letters are
/// opcodes.
///
/// # Instructions
///
/// 'A' <x> <y> <z>: <x> := <y> + <z>
/// 'X' <x> <y> <z>: <x> := <y> * <z>
/// 'U' <x> <y> <z>: <x> := <y> - <z>
/// 'M' <x> <y>: <x> := <y>
/// 'C' <x> <y> <z>: <x> := <y> <= <z> ? 1 : 0 # <x> is a 1 bit flag
/// 'E' <x> <y> <z>: <x> := <y> == <z> ? 1 : 0 # <x> is a 1 bit flag
/// 'B' <x> <y>: branch to <y> if flag <y> is 1, fall-thru otherwise
/// 'J' <x>: jump to
/// 'G' <x>: call to <x>
/// 'F': call to random uuid
/// 'H' <x>: call to stub x
/// 'R': return
/// 'S' <x> <y>: *x := y
/// 'L' <x> <y>: x := *y
#[derive(Debug,Clone)]
pub struct TestArch;

impl TestArch {
    fn read_value(reg: &Region, entry: &mut u64, len: usize, names: &mut Names) -> Result<Value> {
        match reg.read_integer(*entry,Endianess::Little,1).map(|x| x as u8) {
            Ok(b) => match b {
                b'a'...b'z' => {
                    let nam = names.insert(&Name::new(format!("{}",char::from(b)).into(),None));

                    *entry += 1;
                    Value::var(nam,len)
                }
                b'0'...b'9' => {
                    Value::val(Self::read_address(reg,entry)?,32)
                }
                _ => Err(format!("'{}' is nor a variable name neither a constant",b).into())
            }
            _ => Err(format!("premature end while decoding rvalue at {}",*entry).into())
        }
    }

    fn read_variable(reg: &Region, entry: &mut u64, len: usize, names: &mut Names) -> Result<Variable> {
        match reg.read_integer(*entry,Endianess::Little,1).map(|x| x as u8) {
            Ok(b) => match b {
                b'a'...b'z' => {
                    let nam = names.insert(&Name::new(format!("{}",char::from(b)).into(),None));

                    *entry += 1;
                    Variable::new(nam,len)
                }
                _ => Err(format!("'{}' is not a variable name",b).into())
            }
            _ => Err(format!("Premature end while decoding lvalue at {}",*entry).into())
        }
    }

    fn read_address(reg: &Region, entry: &mut u64) -> Result<u64> {
        match reg.read_integer(*entry,Endianess::Little,1).map(|x| x as u8) {
            Ok(b@b'0'...b'9') => {
                let value = (b - b'0') as u64;

                *entry += 1;
                match Self::read_address(reg,entry) {
                    Ok(x) => Ok(value * 10 + x),
                    Err(_) => Ok(value)
                }
            }
            Ok(b) => Err(format!("'{}' is not a number",b).into()),
            _ => Err(format!("Premature end while decoding address at {}",*entry).into())
        }
    }
}

impl Architecture for TestArch {
    type Configuration = ();

    fn prepare(_: &Region, _: &()) -> Result<Cow<'static, [(&'static str, u64, &'static str)]>> {
        Ok(Cow::default())
    }

    fn decode(reg: &Region, mut entry: u64, _: &(), names: &mut Names,
              segments: &mut Segments, strings: &mut Strings,
              matches: &mut Vec<Match>) -> Result<()>
    {
        use Uuid;

        let start = entry;
        let opcode = reg.read_integer(entry,Endianess::Little,1);
        entry += 1;

        match opcode.map(|x| x as u8) {
            Ok(b'M') => {
                let var = Self::read_variable(reg, &mut entry, 32, names)?;
                let val = Self::read_value(reg, &mut entry, 32, names)?;
                let instr = Statement::Expression{
                    result: var.clone(),
                    op: Operation::Move(val.clone())
                };
                let m = Match{
                    area: start..entry,
                    opcode: "mov".into(),
                    operands: SmallVec::from_vec(vec![var.into(), val]),
                    instructions: vec![instr],
                    jumps: SmallVec::from_vec(vec![(start,Value::val(entry, 32)?,Guard::True)]),
                };

                matches.push(m);
                Ok(())
            }
            Ok(b'G') => {
                let addr = Self::read_variable(reg, &mut entry, 32, names)?;
                let instr = Statement::Flow{
                    op: FlowOperation::IndirectCall{ target: addr },
                };
                let m = Match{
                    area: start..entry,
                    opcode: "call".into(),
                    operands: SmallVec::from_vec(vec![addr.into()]),
                    instructions: vec![instr],
                    jumps: SmallVec::from_vec(vec![(start,Value::val(entry, 32)?,Guard::True)]),
                };

                matches.push(m);
                Ok(())
            }
            Ok(b'F') => {
                let instr = Statement::Flow{
                    op: FlowOperation::Call{ function: Uuid::now() },
                };
                let m = Match{
                    area: start..entry,
                    opcode: "call".into(),
                    operands: SmallVec::default(),
                    instructions: vec![instr],
                    jumps: SmallVec::from_vec(vec![(start,Value::val(entry, 32)?,Guard::True)]),
                };

                matches.push(m);
                Ok(())
            }
            Ok(b'H') => {
                let s = match reg.read_integer(entry,Endianess::Little,1) {
                    Ok(b) => {
                        entry += 1;
                        strings.insert(&format!("{}", char::from(b as u8)).into())
                    }
                    _ => {
                        return Err(format!("premature end while decoding stub name {}",entry).into());
                    }
                };
                let instr = Statement::Flow{
                    op: FlowOperation::ExternalCall{ external: s.into() },
                };
                let m = Match{
                    area: start..entry,
                    opcode: "call".into(),
                    operands: SmallVec::default(),
                    instructions: vec![instr],
                    jumps: SmallVec::from_vec(vec![(start,Value::val(entry, 32)?,Guard::True)]),
                };

                matches.push(m);
                Ok(())
            }
            Ok(b'A') => {
                let var = Self::read_variable(reg, &mut entry, 32, names)?;
                let val1 = Self::read_value(reg, &mut entry, 32, names)?;
                let val2 = Self::read_value(reg, &mut entry, 32, names)?;
                let instr = Statement::Expression{
                    result: var,
                    op: Operation::Add(val1.clone(),val2.clone())
                };
                let m = Match{
                    area: start..entry,
                    opcode: "add".into(),
                    operands: SmallVec::from_vec(vec![var.clone().into(),val1.clone(),val2.clone()]),
                    instructions: vec![instr],
                    jumps: SmallVec::from_vec(vec![(start,Value::val(entry, 32)?,Guard::True)]),
                };

                matches.push(m);
                Ok(())
            }
            Ok(b'U') => {
                let var = Self::read_variable(reg, &mut entry, 32, names)?;
                let val1 = Self::read_value(reg, &mut entry, 32, names)?;
                let val2 = Self::read_value(reg, &mut entry, 32, names)?;
                let instr = Statement::Expression{
                    result: var,
                    op: Operation::Subtract(val1.clone(),val2.clone())
                };
                let m = Match{
                    area: start..entry,
                    opcode: "sub".into(),
                    operands: SmallVec::from_vec(vec![var.clone().into(),val1.clone(),val2.clone()]),
                    instructions: vec![instr],
                    jumps: SmallVec::from_vec(vec![(start,Value::val(entry, 32)?,Guard::True)]),
                };

                matches.push(m);
                Ok(())
            }
            Ok(b'X') => {
                let var = Self::read_variable(reg, &mut entry, 32, names)?;
                let val1 = Self::read_value(reg, &mut entry, 32, names)?;
                let val2 = Self::read_value(reg, &mut entry, 32, names)?;
                let instr = Statement::Expression{
                    result: var,
                    op: Operation::Multiply(val1,val2)
                };
                let m = Match{
                    area: start..entry,
                    opcode: "mul".into(),
                    operands: SmallVec::from_vec(vec![var.clone().into(),val1.clone(),val2.clone()]),
                    instructions: vec![instr],
                    jumps: SmallVec::from_vec(vec![(start,Value::val(entry, 32)?,Guard::True)]),
                };

                matches.push(m);
                Ok(())
            }
            Ok(b'C') => {
                let var = Self::read_variable(reg, &mut entry, 1, names)?;
                let val1 = Self::read_value(reg, &mut entry, 32, names)?;
                let val2 = Self::read_value(reg, &mut entry, 32, names)?;
                let instr = Statement::Expression{
                    result: var,
                    op: Operation::LessOrEqualUnsigned(val1,val2)
                };
                let m = Match{
                    area: start..entry,
                    opcode: "leq".into(),
                    operands: SmallVec::from_vec(vec![var.clone().into(),val1.clone(),val2.clone()]),
                    instructions: vec![instr],
                    jumps: SmallVec::from_vec(vec![(start,Value::val(entry, 32)?,Guard::True)]),
                };

                matches.push(m);
                Ok(())
            }
            Ok(b'E') => {
                let var = Self::read_variable(reg, &mut entry, 1, names)?;
                let val1 = Self::read_value(reg, &mut entry, 32, names)?;
                let val2 = Self::read_value(reg, &mut entry, 32, names)?;
                let instr = Statement::Expression{
                    result: var,
                    op: Operation::Equal(val1,val2)
                };
                let m = Match{
                    area: start..entry,
                    opcode: "eq".into(),
                    operands: SmallVec::from_vec(vec![var.clone().into(),val1.clone(),val2.clone()]),
                    instructions: vec![instr],
                    jumps: SmallVec::from_vec(vec![(start,Value::val(entry, 32)?,Guard::True)]),
                };

                matches.push(m);
                Ok(())
            }
            Ok(b'B') => {
                let bit = Self::read_variable(reg, &mut entry, 1, names)?;
                let brtgt = Self::read_address(reg, &mut entry)?;
                let guard = Guard::Predicate{ flag: bit, expected: true };
                let nguard = Guard::Predicate{ flag: bit, expected: false };
                let m = Match{
                    area: start..entry,
                    opcode: "br".into(),
                    operands: SmallVec::from_vec(vec![Value::Variable(bit.clone()),Value::val(brtgt, 32)?]),
                    instructions: vec![],
                    jumps: SmallVec::from_vec(vec![
                        (start,Value::val(entry, 32)?,nguard),
                        (start,Value::val(brtgt, 32)?,guard)
                    ]),
                };

                matches.push(m);
                Ok(())
            }
            Ok(b'J') => {
                let jtgt = Self::read_value(reg, &mut entry, 32, names)?;
                let m = Match{
                    area: start..entry,
                    opcode: "jmp".into(),
                    operands: SmallVec::from_vec(vec![jtgt]),
                    instructions: vec![],
                    jumps: SmallVec::from_vec(vec![(start,jtgt,Guard::True)]),
                };

                matches.push(m);
                Ok(())
            }
            Ok(b'R') => {
                let instr = Statement::Flow{
                    op: FlowOperation::Return
                };
                let m = Match{
                    area: start..entry,
                    opcode: "ret".into(),
                    operands: SmallVec::default(),
                    instructions: vec![instr],
                    jumps: SmallVec::default(),
                };

                matches.push(m);
                Ok(())
            }
            Ok(b'S') => {
                let addr = Self::read_value(reg, &mut entry, 32, names)?;
                let val = Self::read_value(reg, &mut entry, 32, names)?;
                let seg = Segment{ name: segments.insert(&Name::new("ram".into(),None)) };
                let instr = Statement::Memory{
                    result: seg,
                    op: MemoryOperation::Store{
                        segment: seg,
                        endianess: Endianess::Little,
                        bytes: 4,
                        address: addr,
                        value: val,
                    }
                };
                let m = Match{
                    area: start..entry,
                    opcode: "store".into(),
                    operands: SmallVec::from_vec(vec![addr.clone().into(),val.clone()]),
                    instructions: vec![instr],
                    jumps: SmallVec::from_vec(vec![(start,Value::val(entry, 32)?,Guard::True)]),
                };

                matches.push(m);
                Ok(())
            }
            Ok(b'L') => {
                let val = Self::read_variable(reg, &mut entry, 32, names)?;
                let addr = Self::read_value(reg, &mut entry, 32, names)?;
                let seg = segments.insert(&Name::new("ram".into(),None));
                let instr = Statement::Expression{
                    result: val,
                    op: Operation::Load(Segment{ name: seg },Endianess::Little,4,addr)
                };
                let m = Match{
                    area: start..entry,
                    opcode: "load".into(),
                    operands: SmallVec::from_vec(vec![val.clone().into(),addr.clone()]),
                    instructions: vec![instr],
                    jumps: SmallVec::from_vec(vec![(start,Value::val(entry, 32)?,Guard::True)]),
                };

                matches.push(m);
                Ok(())
            }
            Ok(o) => {
                Err(format!("Unknown opcode '{}' at {}",o,start).into())
            }
            _ => {
                Err(format!("Premature end while decoding opcode {}",start).into())
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use {Function, Region, Uuid};

    #[test]
    fn test_arch() {
        let reg = Region::from_buf("",16,b"Mi1MiiAiiiAi1iAii1Ai11CiiiCi1iCii1Ci11Bi0RSiiS1iS11LiiLi1".to_vec(),0);
        let _ = Function::new::<TestArch>((), 0, &reg, Uuid::now()).unwrap();
    }
}
