#define _XOPEN_SOURCE 600

#include "pentest_state.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>





void pentest_state_init(struct pentest_state *p_s, size_t max_targets,
						size_t max_attacks, size_t max_scans)
{
    p_s->max_targets = max_targets;
    p_s->num_targets = 0;
    p_s->targets = malloc(max_targets * sizeof(char*));

    p_s->max_attacks = max_attacks;
    p_s->num_attacks = 0;
    p_s->attacks = malloc(max_attacks *sizeof(struct active_attack*));

	p_s->max_scans = max_scans;
    p_s->num_scans = 0;
    p_s->scans = malloc(max_scans *sizeof(struct active_scan*));



    p_s->max_failed = PENTEST_STATE_INIT_NUM_PAST_ATTACKS;
    p_s->num_failed = 0;
    p_s->failed = malloc(p_s->max_failed * sizeof(struct past_attack));

    p_s->max_worked = PENTEST_STATE_INIT_NUM_PAST_ATTACKS;
    p_s->num_worked = 0;
    p_s->worked = malloc(p_s->max_worked * sizeof(struct past_attack));

	p_s->max_done = PENTEST_STATE_INIT_NUM_PAST_SCANS;
    p_s->num_done = 0;
    p_s->done = malloc(p_s->max_done * sizeof(struct past_scan));

    pthread_mutex_init(&p_s->lock,NULL);

    
}

void pentest_state_cleanup(struct pentest_state *p_s)
{
    
    for(size_t i = 0;i < p_s->num_targets;i++)
    {
		free(p_s->targets[i]);
    }
    free(p_s->targets);
    
    //destroy mutex too
    for(size_t i=0; i < p_s->num_failed ; i++)
    {
		past_attack_cleanup(&p_s->failed[i]);
    }
    free(p_s->failed);

    for(size_t i=0; i < p_s->num_worked ; i++)
    {
		past_attack_cleanup(&p_s->worked[i]);
    }
    free(p_s->worked);

    for(size_t i=0; i<p_s->num_attacks;i++)
    {
		active_attack_cleanup(p_s->attacks[i]);
    }
    free(p_s->attacks);

    pthread_mutex_destroy(&p_s->lock);
}


void pentest_state_add_target(struct pentest_state *p_s, char *target)
{
    pthread_mutex_lock(&p_s->lock);
    if(p_s->num_targets >= p_s->max_targets)
    {
		fprintf(stderr,"[Panic] Trying to add a target beyond the allowed threshold\n");
    }

    p_s->targets[p_s->num_targets] = strdup(target);
    p_s->num_targets++;
    pthread_mutex_unlock(&p_s->lock);
}


void pentest_state_add_attack(struct pentest_state *p_s, struct active_attack *attack)
{
    pthread_mutex_lock(&p_s->lock);
    if(p_s->num_attacks >= p_s->max_attacks)
    {
		fprintf(stderr,"[Panic] Trying to add an attack beyond the allowed threshold\n");
    }

    p_s->attacks[p_s->num_attacks] = attack;
    p_s->num_attacks++;
    pthread_mutex_unlock(&p_s->lock);
}
void pentest_state_add_scan(struct pentest_state *p_s, struct active_scan *scan)
{
	pthread_mutex_lock(&p_s->lock);
    if(p_s->num_scans >= p_s->max_scans)
    {
		fprintf(stderr,"[Panic] Trying to add an scan beyond the allowed threshold\n");
    }

    p_s->scans[p_s->num_scans] = scan;
    p_s->num_scans++;
    pthread_mutex_unlock(&p_s->lock);
}



//Do we also do memory cleanup here?
void pentest_state_rem_attack(struct pentest_state *p_s, struct active_attack *a)
{
    pthread_mutex_lock(&p_s->lock);

    for(size_t i = 0; i < p_s->num_attacks ; i++)
    {
		struct active_attack *a_a = p_s->attacks[i];
		if(!strcmp(a->attack->name,a_a->attack->name) && !strcmp(a->target,a_a->target))
		{
			if(i != p_s->num_attacks -1 )
			{
				p_s->attacks[i] = p_s->attacks[p_s->num_attacks-1];
			}
			p_s->num_attacks--;

			active_attack_cleanup(a_a);
	    
			break;
		}
    }
    pthread_mutex_unlock(&p_s->lock);
    
}

void pentest_state_rem_scan(struct pentest_state *p_s, struct active_scan *scan)
{
	pthread_mutex_lock(&p_s->lock);

    for(size_t i = 0; i < p_s->num_scans ; i++)
    {
		struct active_scan *a_s = p_s->scans[i];
		if(active_scan_eq(scan,a_s))
		{
			if(i != p_s->num_scans -1 )
			{
				p_s->scans[i] = p_s->scans[p_s->num_scans-1];
			}
			p_s->num_scans--;

			active_scan_cleanup(a_s);
	    
			break;
		}
    }
    pthread_mutex_unlock(&p_s->lock);
    
}



void pentest_state_rem_target(struct pentest_state *p_s, char *target)
{
    //pthread_mutex_lock(&p_s->lock);

    for(size_t i = 0; i < p_s->num_targets ; i++)
    {
		char *a_target = p_s->targets[i];
		if(!strcmp(a_target,target))
		{
			if(i != p_s->num_targets -1 )
			{
				p_s->targets[i] = p_s->targets[p_s->num_targets-1];
			}
			p_s->num_targets--;

			//Free target?
	    
			break;
		}
    }
    pthread_mutex_unlock(&p_s->lock);
}



void pentest_state_add_failed(struct pentest_state *p_s, char *a_p, char *target)
{
    pthread_mutex_lock(&p_s->lock);
    if(p_s->num_failed >= p_s->max_failed)
    {
		//can fail spectacularly if malloc fails
		p_s->max_failed *= 2;
		p_s->failed = realloc(p_s->failed, p_s->max_failed*sizeof(struct past_attack));
    }

    past_attack_init(&p_s->failed[p_s->num_failed], a_p, target);
    p_s->num_failed++;
    pthread_mutex_unlock(&p_s->lock);
}

void pentest_state_add_worked(struct pentest_state *p_s, char *a_p, char *target)
{
    pthread_mutex_lock(&p_s->lock);
    if(p_s->num_worked >= p_s->max_worked)
    {
		//can fail spectacularly if malloc fails
		p_s->max_worked *= 2;
		p_s->worked = realloc(p_s->worked, p_s->max_worked*sizeof(struct past_attack));
    }

    past_attack_init(&p_s->worked[p_s->num_worked], a_p, target);
    p_s->num_worked++;
    pthread_mutex_unlock(&p_s->lock);
}


void pentest_state_add_done(struct pentest_state *p_s, char *scan, char *target)
{
	pthread_mutex_lock(&p_s->lock);
    if(p_s->num_done >= p_s->max_done)
    {
		//can fail spectacularly if malloc fails
		p_s->max_done *= 2;
		p_s->done = realloc(p_s->done, p_s->max_done*sizeof(struct past_scan));
    }

    past_scan_init(&p_s->done[p_s->num_done], scan, target);
    p_s->num_done++;
    pthread_mutex_unlock(&p_s->lock);
}


    
bool pentest_active_target(struct pentest_state *p_s, char *target)
{
    bool res = false;

    pthread_mutex_lock(&p_s->lock);
    
    for(size_t i = 0; i < p_s->num_targets;i++)
    {
		if(strcmp(p_s->targets[i],target) == 0)
		{
			res = true;
			break;
		}
    }

    pthread_mutex_unlock(&p_s->lock);
    
    return res;
}


bool pentest_exploited_target(struct pentest_state *p_s, char *target)
{
    bool res = false;
    pthread_mutex_lock(&p_s->lock);
    
    pthread_mutex_unlock(&p_s->lock);
    return res;
}


bool pentest_attack_active(struct pentest_state *p_s, char *attack, char *target)
{

    bool res = false;
    pthread_mutex_lock(&p_s->lock);
    for(size_t i = 0; i < p_s->num_attacks;i++)
    {
		struct attack_pattern *a_p = p_s->attacks[i]->attack;
		char *a_target = p_s->attacks[i]->target;
		if(strcmp(a_p->name, attack) ==0 && strcmp(a_target,target) == 0)
		{
			res = true;
			break;
		}
    }

    pthread_mutex_unlock(&p_s->lock);
    
    return res;
}

bool pentest_attack_failed(struct pentest_state *p_s, char *attack, char *target)
{

    bool res = false;
    pthread_mutex_lock(&p_s->lock);
    for(size_t i = 0; i < p_s->num_failed; i++)
    {
		struct past_attack *p_a = &p_s->failed[i];
		if(!strcmp(p_a->a_p, attack) && !strcmp(p_a->target, target))
		{
			res = true;
			break;
		}
    }

    pthread_mutex_unlock(&p_s->lock);
    return res;
}

bool pentest_attack_worked(struct pentest_state *p_s, char *attack, char *target)
{
    bool res = false;
    pthread_mutex_lock(&p_s->lock);
    for(size_t i = 0; i < p_s->num_worked; i++)
    {
		struct past_attack *p_a = &p_s->worked[i];
		if(!strcmp(p_a->a_p, attack) && !strcmp(p_a->target, target))
		{
			res = true;
			break;
		}
    }

    pthread_mutex_unlock(&p_s->lock);
    return res;
}


size_t pentest_state_num_attacks(struct pentest_state *p_s)
{
    size_t num_attacks = 0;
    pthread_mutex_lock(&p_s->lock);
    num_attacks = p_s->num_attacks;
    pthread_mutex_unlock(&p_s->lock);

    return num_attacks;
    
}
size_t pentest_state_max_attacks(struct pentest_state *p_s)
{
    size_t max_attacks = 0;
    pthread_mutex_lock(&p_s->lock);
    max_attacks = p_s->max_attacks;
    pthread_mutex_unlock(&p_s->lock);

    return max_attacks;
    
}
size_t pentest_state_num_targets(struct pentest_state *p_s)
{
    size_t num_targets = 0;
    pthread_mutex_lock(&p_s->lock);
    num_targets = p_s->num_targets;
    pthread_mutex_unlock(&p_s->lock);

    return num_targets;
}
size_t pentest_state_max_targets(struct pentest_state *p_s)
{
    size_t max_targets = 0;
    pthread_mutex_lock(&p_s->lock);
    max_targets = p_s->max_targets;
    pthread_mutex_unlock(&p_s->lock);

    return max_targets;
}
size_t pentest_state_num_scans(struct pentest_state *p_s)
{
    size_t num_scans = 0;
    pthread_mutex_lock(&p_s->lock);
    num_scans = p_s->num_scans;
    pthread_mutex_unlock(&p_s->lock);

    return num_scans;
}
size_t pentest_state_max_scans(struct pentest_state *p_s)
{
    size_t max_scans = 0;
    pthread_mutex_lock(&p_s->lock);
    max_scans = p_s->max_scans;
    pthread_mutex_unlock(&p_s->lock);

    return max_scans;
}

size_t pentest_state_num_failed(struct pentest_state *p_s)
{
    size_t num_failed = 0;
    pthread_mutex_lock(&p_s->lock);
    num_failed = p_s->num_failed;
    pthread_mutex_unlock(&p_s->lock);

    return num_failed;
}
size_t pentest_state_max_failed(struct pentest_state *p_s)
{
    size_t max_failed = 0;
    pthread_mutex_lock(&p_s->lock);
    max_failed = p_s->max_failed;
    pthread_mutex_unlock(&p_s->lock);

    return max_failed;
}
size_t pentest_state_num_worked(struct pentest_state *p_s)
{
    size_t num_worked = 0;
    pthread_mutex_lock(&p_s->lock);
    num_worked = p_s->num_worked;
    pthread_mutex_unlock(&p_s->lock);

    return num_worked;
}


size_t pentest_state_max_worked(struct pentest_state *p_s)
{
    size_t max_worked = 0;
    pthread_mutex_lock(&p_s->lock);
    max_worked = p_s->max_worked;
    pthread_mutex_unlock(&p_s->lock);

    return max_worked;
}

bool pentest_scan_active(struct pentest_state *p_s, char *scan, char *target)
{
 bool res = false;
    pthread_mutex_lock(&p_s->lock);
    for(size_t i = 0; i < p_s->num_scans;i++)
    {
		struct scan *a_s = p_s->scans[i]->scan;
		char *a_target = p_s->scans[i]->target;

		if((!target && a_target) || (target && !a_target))
		{
			continue;
		}
		
		if(strcmp(a_s->name, scan) ==0 &&
		   ( ((!a_target && !target) || strcmp(a_target,target) == 0)))
		{
			res = true;
			break;
		}
    }

    pthread_mutex_unlock(&p_s->lock);
    
    return res;
	
}


bool pentest_scan_done(struct pentest_state *p_s, char *scan, char *target)
{

  bool res = false;
  pthread_mutex_lock(&p_s->lock);
  for(size_t i = 0; i < p_s->num_done; i++)
  {
	  struct past_scan *s = &p_s->done[i];

	  if( (!target && s->target) || (target && !s->target))
	  {
		  continue;
	  }
	  
	  if(!strcmp(s->scan, scan) && ((!s->target && !target)  || !strcmp(s->target, target) ))
	  {
		  res = true;
		  break;
	  }
  }

  pthread_mutex_unlock(&p_s->lock);
  return res;

}



char *pentest_state_get_target(struct pentest_state *p_s, size_t idx)
{
    char *target = NULL;

    pthread_mutex_lock(&p_s->lock);
    if(idx < p_s->num_targets)
    {
		target = p_s->targets[idx];
    }
    pthread_mutex_unlock(&p_s->lock);

    return target;
}

struct active_attack *pentest_state_get_attack(struct pentest_state *p_s, size_t idx)
{
    struct active_attack *a_a = NULL;

    pthread_mutex_lock(&p_s->lock);
    if(idx < p_s->num_attacks)
    {
		a_a = p_s->attacks[idx];
    }
    pthread_mutex_unlock(&p_s->lock);

    return a_a;
}

struct active_scan *pentest_state_get_scan(struct pentest_state *p_s, size_t idx)
{
	struct active_scan *a_s = NULL;

    pthread_mutex_lock(&p_s->lock);
    if(idx < p_s->num_scans)
    {
		a_s = p_s->scans[idx];
    }
    pthread_mutex_unlock(&p_s->lock);

    return a_s;

}


void pentest_state_active_attack_update(struct pentest_state *p_s, struct active_attack *a_a)
{

    //There's probably a better way to do this.
    if(a_a->cur_step_state != AAS_DONE)
    {
		return;
    }
    
    //current step done
    if(a_a->cur_step_state == AAS_DONE)
    {
		//Did we succeed?)
		if(a_a->cur_res)
		{
			if(a_a->cur_step == a_a->attack->num_steps -1)
			{
				//all steps complete, move to the attacks that worked and remove from active
				pentest_state_add_worked(p_s,a_a->attack->name, a_a->target);
				pentest_state_rem_attack(p_s,a_a);
		
			}
			else
			{
				//move on to the next step
				active_attack_next_step(a_a);
			}
		}
		else
		{
			//If a step fails, the whole attack fails
			pentest_state_add_failed(p_s,a_a->attack->name, a_a->target);
			pentest_state_rem_attack(p_s,a_a);
		}
    }
    
}

void pentest_state_active_scan_update(struct pentest_state *p_s, struct active_scan *a_s)
{
	//current step done
    if(a_s->scan_state == AS_DONE)
    {
		//scan complete, move done nad remove from active
		pentest_state_add_done(p_s,a_s->scan->name, a_s->target);
		pentest_state_rem_scan(p_s,a_s);		
    }
    
}