/*
 * See the dyninst/COPYRIGHT file for copyright information.
 * 
 * We provide the Paradyn Tools (below described as "Paradyn")
 * on an AS IS basis, and do not warrant its validity or performance.
 * We reserve the right to update, modify, or discontinue this
 * software at any time.  We shall have no obligation to supply such
 * updates or modifications or any other form of support to you.
 * 
 * By your use of Paradyn, you understand and agree that we (or any
 * other person or entity with proprietary rights in Paradyn) are
 * under no obligation to provide either maintenance services,
 * update services, notices of latent defects, or correction of
 * defects for Paradyn.
 * 
 * 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
 */
#include <algorithm>

#include "dyntypes.h"

#include "CodeObject.h"
#include "CFG.h"

#include "Parser.h"
#include "debug_parse.h"
#include "util.h"

#include "dataflowAPI/h/slicing.h"
#include "dataflowAPI/h/AbslocInterface.h"
#include "instructionAPI/h/InstructionDecoder.h"
#include "dynutil/h/Graph.h"
#include "StackTamperVisitor.h"

using namespace std;

using namespace Dyninst;
using namespace Dyninst::ParseAPI;

Function::Function() :
        _start(0),
        _obj(NULL),
        _region(NULL),
        _isrc(NULL),
        _src(RT),
        _rs(UNSET),
        _entry(NULL),
	 _is_leaf_function(true),
	 _ret_addr(0),
        _parsed(false),
        _cache_valid(false),
        _no_stack_frame(true),
        _saves_fp(false),
        _cleans_stack(false),
        _tamper(TAMPER_UNSET),
        _tamper_addr(0)
{
    fprintf(stderr,"PROBABLE ERROR, default ParseAPI::Function constructor\n");
}

Function::Function(Address addr, string name, CodeObject * obj, 
    CodeRegion * region, InstructionSource * isrc) :
        _start(addr),
        _obj(obj),
        _region(region),
        _isrc(isrc),
        _src(RT),
        _rs(UNSET),
        _name(name),
        _entry(NULL),
	 _is_leaf_function(true),
	 _ret_addr(0),
        _parsed(false),
        _cache_valid(false),
        _no_stack_frame(true),
        _saves_fp(false),
        _cleans_stack(false),
        _tamper(TAMPER_UNSET),
        _tamper_addr(0)
{
    if (obj->defensiveMode()) {
        mal_printf("new funct at %lx\n",addr);
    }
    if (obj && obj->cs()) {
        obj->cs()->incrementCounter(PARSE_FUNCTION_COUNT);
    }
}

ParseAPI::Edge::~Edge() {
}

Function::~Function()
{
    if (_obj && _obj->cs()) {
        _obj->cs()->decrementCounter(PARSE_FUNCTION_COUNT);
    }
    vector<FuncExtent *>::iterator eit = _extents.begin();
    for( ; eit != _extents.end(); ++eit) {
        delete *eit;
    }
}

Function::blocklist &
Function::blocks()
{
    if(!_cache_valid)
        finalize();
    return _bl;
}

// Get the current set of blocks,
// as a const operation
const Function::blocklist&
Function::blocks() const
{
  /*Function* mutable_this = const_cast<Function*>(this);
  
  if(!_cache_valid)
    mutable_this->finalize();
  */
  assert(_cache_valid);
  
  return _bl;
}


const Function::edgelist & 
Function::callEdges() {
    if(!_cache_valid)
        finalize();
    return _call_edge_list; 
}

const Function::blocklist &
Function::returnBlocks() {
  if (!_cache_valid) 
    finalize();
  return _retBL;
}

const Function::blocklist &
Function::exitBlocks() {
  if (!_cache_valid) 
    finalize();
  return _exitBL;
}

vector<FuncExtent *> const&
Function::extents()
{
    if(!_cache_valid)
        finalize(); 
    return _extents;
}

void
Function::finalize()
{
  _extents.clear();
  _exitBL.clear();
  _bl.clear();
  _retBL.clear();
  _call_edge_list.clear();

    // The Parser knows how to finalize
    // a Function's parse data
    _obj->parser->finalize(this);
}

vector<Block *> const&
Function::blocks_int()
{
    if(_cache_valid || !_entry)
        return _bl;

    // overloaded map warning:
    // visited[addr] == 1 means visited
    // visited[addr] == 2 means already on the return list
    dyn_hash_map<Address,short> visited;
    vector<Block *> worklist;

    bool need_entry = true;
    for(vector<Block*>::iterator bit=_bl.begin();
        bit!=_bl.end();++bit) 
    {
        Block * b = *bit;
        visited[b->start()] = 1;
        need_entry = need_entry && (b != _entry);
    }
    worklist.insert(worklist.begin(),_bl.begin(),_bl.end());

    if(need_entry) {
        worklist.push_back(_entry);
        visited[_entry->start()] = 1;
        add_block(_entry);
    }

    // avoid adding duplicate return blocks
    for(vector<Block*>::iterator bit=_exitBL.begin();
        bit!=_exitBL.end();++bit)
    {
        Block * b = *bit;
        visited[b->start()] = 2;
    }
    
    while(!worklist.empty()) {
        Block * cur = worklist.back();
        worklist.pop_back();

        bool link_return = false;
        bool exits_func = false;
        bool found_call = false;
        bool found_call_ft = false;
        const Block::edgelist & trgs = cur->targets();
        if (trgs.empty()) {
           // Woo hlt!
	  parsing_printf("No targets, exits func\n");
           exits_func = true;
        }
        for(Block::edgelist::const_iterator tit=trgs.begin();
            tit!=trgs.end();++tit) {
            Edge * e = *tit;
            Block * t = e->trg();

            parsing_printf("\t Considering target block 0x%lx from edge %p\n", t->start(), e); 

            if (e->type() == CALL_FT) {
               found_call_ft = true;
            }

            if(e->type() == CALL) {
               parsing_printf("\t Call typed\n");
                _call_edge_list.insert(e);
                found_call = true;
                continue;
            }

            if(e->type() == RET) {
                link_return = true;
                exits_func = true;
		parsing_printf("Block has return edge\n");
                if (obj()->defensiveMode()) {
                    if (_tamper != TAMPER_UNSET && _tamper != TAMPER_NONE) 
                       continue;
                }
                
                set_retstatus(RETURN);
                continue;
            }

            /* Handle tailcall edges */
            if(e->interproc()) {
               parsing_printf("\t Interprocedural\n");
                _call_edge_list.insert(e);
                found_call = true;
                continue;
            }
            // If we are heading to a different CodeObject, call it a return
            // and don't add target blocks.
            if (t->obj() != cur->obj()) {
               // This is a jump to a different CodeObject; call it an exit
	      parsing_printf("Block exits object\n");
               exits_func = true;
                continue;
            }

            /* sink edges receive no further processing */
            if(e->sinkEdge()) {
               parsing_printf("\t Sink edge, skipping\n");
                continue;
            }

            if(!HASHDEF(visited,t->start())) {
               parsing_printf("\t Adding target block to worklist\n");
                worklist.push_back(t);
                visited[t->start()] = true;
                add_block(t);
            }
        }
        if (found_call && !found_call_ft && !obj()->defensiveMode()) {
	  parsing_printf("\t exits func\n");
           exits_func = true;
        }

        if (link_return) assert(exits_func);

        if(exits_func) {
           if (link_return)
              delayed_link_return(_obj,cur);
           if(visited[cur->start()] <= 1) {
              _exitBL.push_back(cur);
	      parsing_printf("Adding block 0x%lx as exit\n", cur->start());
              if (link_return) 
	      {
                 _retBL.push_back(cur);
	      }
	      
           }
        }
    }
    Block::compare comp;
    sort(_bl.begin(),_bl.end(),comp);

    return _bl;
}

/* Adds return edges to the CFG for a particular retblk, based 
 * on callers to this function.  Handles case of return block 
 * that targets the entry block of a new function separately
 * (ret to entry happens if the function tampers with its stack 
 *  and maybe if this function is a signal handler?) 
 */ 
void
Function::delayed_link_return(CodeObject * o, Block * retblk)
{
    bool link_entry = false;

    dyn_hash_map<Address,bool> linked;
    Block::edgelist::const_iterator eit = retblk->targets().begin();
    for( ; eit != retblk->targets().end(); ++eit) {
        Edge * e = *eit;
        linked[e->trg()->start()] = true;
    }

    eit = _entry->sources().begin();
    for( ; eit != _entry->sources().end(); ++eit) {
        Edge * e = *eit;
        if(e->type() == CALL) {
            parsing_printf("[%s:%d] linking return edge %lx -> %lx\n",
                FILE__,__LINE__,retblk->lastInsnAddr(),e->src()->end());

            // XXX opportunity here to be more conservative about delayed
            //     determination of return status
    
            Block * call_ft = _obj->findBlockByEntry(region(),e->src()->end());
            if(!call_ft) {
                parsing_printf("[%s:%d] no block found, error!\n",
                    FILE__,__LINE__);
            } 
            else if(!HASHDEF(linked,call_ft->start())) {
                if(call_ft == _entry)
                    link_entry = true;
                else 
                    o->add_edge(retblk,call_ft,RET);
                linked[call_ft->start()] = true;
            }
        }
    }
    // can't do this during iteration
    if(link_entry)
        o->add_edge(retblk,_entry,RET);
}

void
Function::add_block(Block *b)
{
  ++b->_func_cnt;            // block counts references
  _bl.push_back(b);
  _bmap[b->start()] = b;
}

const string &
Function::name() 
{
    return _name;
}

bool
Function::contains(Block *b)
{
    if(!_cache_valid)
        finalize();

    return HASHDEF(_bmap,b->start());
}

void Function::setEntryBlock(Block *new_entry)
{
    obj()->parser->move_func(this, new_entry->start(), new_entry->region());
    _region = new_entry->region();
    _start = new_entry->start();
    _entry = new_entry;
}

void Function::set_retstatus(FuncReturnStatus rs) 
{
    // If we are changing the return status, update prev counter
    if (_rs != UNSET) {
        if (_rs == NORETURN) {
            _obj->cs()->decrementCounter(PARSE_NORETURN_COUNT);
        } else if (_rs == RETURN) {
            _obj->cs()->decrementCounter(PARSE_RETURN_COUNT);
        } else if (_rs == UNKNOWN) {
            _obj->cs()->decrementCounter(PARSE_UNKNOWN_COUNT);
        }
    }

    // Update counter information
    if (rs == NORETURN) {
        _obj->cs()->incrementCounter(PARSE_NORETURN_COUNT);
    } else if (rs == RETURN) {
        _obj->cs()->incrementCounter(PARSE_RETURN_COUNT);
    } else if (rs == UNKNOWN) {
        _obj->cs()->incrementCounter(PARSE_UNKNOWN_COUNT);
    }
    _rs = rs;
}

void 
Function::removeBlock(Block* dead)
{
    _cache_valid = false;
    bool found = false;

    // remove dead block from _bl // KEVINTODO: use binary search
    std::vector<Block *>::iterator biter = _bl.begin();
    while ( !found && _bl.end() != biter ) {
        if (dead == *biter) {
            found = true;
            biter = _bl.erase(biter);
        }
        else {
            biter++;
        }
    }
    if (!found) {
        fprintf(stderr,"Error, tried to remove block [%lx,%lx) from "
                "function at %lx that it does not belong to at %s[%d]\n",
                dead->start(),dead->end(), addr(), FILE__,__LINE__);
        assert(0);
    }

    // specify replacement entry prior to deleting entry block, unless 
    // deleting all blocks
    if (dead == _entry) {
        mal_printf("Warning: removing entry block [%lx %lx) for function at "
                   "%lx\n", dead->start(), dead->end(), addr());
        _entry = NULL;
        assert(0);
    }

    // remove dead block from _retBL and _call_edge_list
    const Block::edgelist & outs = dead->targets();
    for (Block::edgelist::const_iterator oit = outs.begin();
         outs.end() != oit; 
         oit++ ) 
    {
        switch((*oit)->type()) {
            case CALL: {
                bool foundEdge = false;
                for (set<Edge*>::iterator cit = _call_edge_list.begin();
                     _call_edge_list.end() != cit;
                     cit++) 
                {
                    if (*oit == *cit) {
                        foundEdge = true;
                        _call_edge_list.erase(cit);
                        break;
                    }
                }
                assert(foundEdge || (*oit)->sinkEdge());
                break;
            }
            case RET:
                _retBL.erase(std::remove(_retBL.begin(),
                                                 _retBL.end(),
                                                 dead),
                                     _retBL.end());
                break;
            default:
                break;
        }
    }
    // remove dead block from block map
    _bmap.erase(dead->start());
}

class ST_Predicates : public Slicer::Predicates {};

StackTamper 
Function::tampersStack(bool recalculate)
{
    using namespace SymbolicEvaluation;
    using namespace InstructionAPI;

    if ( ! obj()->defensiveMode() ) { 
        assert(0);
        _tamper = TAMPER_NONE;
        return _tamper;
    }

    // this is above the cond'n below b/c it finalizes the function, 
    // which could in turn call this function
    const Function::blocklist retblks(returnBlocks());
    if ( retblks.begin() == retblks.end() ) {
        _tamper = TAMPER_NONE;
        return _tamper;
    }
	_cache_valid = false;

    // if we want to re-calculate the tamper address
    if (!recalculate && TAMPER_UNSET != _tamper) {
        return _tamper;
    }

    AssignmentConverter converter(true);
    vector<Assignment::Ptr> assgns;
    ST_Predicates preds;
    _tamper = TAMPER_UNSET;
    Function::blocklist::const_iterator bit;
    for (bit = retblks.begin(); retblks.end() != bit; ++bit) {
        Address retnAddr = (*bit)->lastInsnAddr();
        InstructionDecoder retdec(this->isrc()->getPtrToInstruction(retnAddr), 
                                  InstructionDecoder::maxInstructionLength, 
                                  this->region()->getArch() );
        Instruction::Ptr retn = retdec.decode();
        converter.convert(retn, retnAddr, this, *bit, assgns);
        vector<Assignment::Ptr>::iterator ait;
        AST::Ptr sliceAtRet;

        for (ait = assgns.begin(); assgns.end() != ait; ait++) {
            AbsRegion & outReg = (*ait)->out();
            if ( outReg.absloc().isPC() ) {
                // First check to see if an input is an unresolved stack slot 
                // (or worse, the heap) - since if that's the case there's no use
                // in spending a lot of time slicing.
                std::vector<AbsRegion>::const_iterator in_iter;
                for (in_iter = (*ait)->inputs().begin();
                    in_iter != (*ait)->inputs().end(); ++in_iter) {
                    if (in_iter->type() != Absloc::Unknown) {
                        _tamper = TAMPER_NONZERO;
                        _tamper_addr = 0;
                        set_retstatus(NORETURN);
                        mal_printf("Stack tamper analysis for ret block at "
                               "%lx found unresolved stack slot or heap "
                               "addr, marking as TAMPER_NONZERO\n", retnAddr);
                        return _tamper;
                    }
                }

                Slicer slicer(*ait,*bit,this);
                Graph::Ptr slGraph = slicer.backwardSlice(preds);
                if (dyn_debug_malware) {
                    stringstream graphDump;
                    graphDump << "sliceDump_" << this->name() << "_" 
                              << hex << retnAddr << dec << ".dot";
                    //slGraph->printDOT(graphDump.str());
                }
                DataflowAPI::Result_t slRes;
                DataflowAPI::SymEval::expand(slGraph,slRes);
                sliceAtRet = slRes[*ait];
                if (dyn_debug_malware && sliceAtRet != NULL) {
                    cerr << "assignment " << (*ait)->format() << " is "
                         << sliceAtRet->format() << "\n";
                }
                break;
            }
        }
        if (sliceAtRet == NULL) {
            mal_printf("Failed to produce a slice for retn at %x %s[%d]\n",
                       retnAddr, FILE__,__LINE__);
            continue;
        } 
        StackTamperVisitor vis(Absloc(-1 * isrc()->getAddressWidth(), 0, this));
        Address curTamperAddr=0;
        StackTamper curtamper = vis.tampersStack(sliceAtRet, curTamperAddr);
        mal_printf("StackTamperVisitor for func at 0x%lx block[%lx %lx) w/ "
                   "lastInsn at 0x%lx returns tamper=%d tamperAddr=0x%lx\n",
                   _start, (*bit)->start(), (*bit)->end(), retnAddr, 
                   curtamper, curTamperAddr);
        if (TAMPER_UNSET == _tamper || TAMPER_NONE == _tamper ||
            (TAMPER_NONZERO == _tamper && 
             TAMPER_NONE != curtamper))
        {
            _tamper = curtamper;
            _tamper_addr = curTamperAddr;
        } 
        else if ((TAMPER_REL == _tamper   || TAMPER_ABS == _tamper) &&
                 (TAMPER_REL == curtamper || TAMPER_ABS == curtamper))
        {
            if (_tamper != curtamper || _tamper_addr != curTamperAddr) {
                fprintf(stderr, "WARNING! Unhandled case in stackTamper "
                        "analysis, func at %lx has distinct tamperAddrs "
                        "%d:%lx %d:%lx at different return instructions, "
                        "setting to TAMPER_NONZERO %s[%d]\n", 
                        this->addr(), _tamper,_tamper_addr, curtamper, 
                        curTamperAddr, FILE__, __LINE__);
                _tamper = TAMPER_NONZERO; // let instrumentation take care of it
            }
        }
        assgns.clear();
    }

    if ( TAMPER_UNSET == _tamper ) {
        mal_printf("WARNING: we found no valid slices for function at %lx "
                   "%s[%d]\n", _start, _tamper_addr, FILE__,__LINE__);
        _tamper = TAMPER_NONZERO;
    }

    if ( TAMPER_NONE != _tamper && TAMPER_REL != _tamper && RETURN == _rs ) {
        set_retstatus(NORETURN);
    }
    return _tamper;
}

void Function::destroy(Function *f) {
   f->obj()->destroy(f);
}
