package uk.co.pentest.psychoPATH;

import burp.BurpExtender;
import burp.IBurpExtenderCallbacks;
import burp.IIntruderPayloadGenerator;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.swing.ListModel;

/**
 *
 * @author ewilded
 */
public final class IntruderPayloadGeneratorByte implements IIntruderPayloadGenerator 
{
    int payloadIndex;    // counter for the payload mark       
    byte[] bytes_raw;    // byte payloads    
    ArrayList<String> bytes_out; // 
    IBurpExtenderCallbacks callbacks = BurpExtender.getBurpCallbacks();
    PsychoTab tab;

    public IntruderPayloadGeneratorByte(PsychoTab tab) 
    {
        this.tab=tab;
        this.tab.psychoPanel.logOutput("Byte is being created.");
        int byteIndex=0;
        this.bytes_out=new ArrayList<>();
            /*
                    "non-alpha";
                    "alpha";
                    "non-alpha-print";
                    "non-alpha-non-print";
                    "all";
            */
            //this.bytes_raw=new ArrayList<>();
        if("non-alpha".equals(this.tab.psychoPanel.byteGeneratorRange))
        {
                this.bytes_raw = new byte[194];
                // 0-47     48
                // 58-64    7
                // 91-96    6
                // 123-255  133
                // 194 all together
                for(int i=0;i<48;i++)
                {
                    bytes_raw[byteIndex]=(byte)i;
                    byteIndex++;
                }
                for(int i=58;i<65;i++)
                {
                    bytes_raw[byteIndex]=(byte)i;
                    byteIndex++;                
                }
                for(int i=91;i<97;i++)
                {
                    bytes_raw[byteIndex]=(byte)i;
                    byteIndex++;                
                }
                for(int i=123;i<256;i++)
                {
                    bytes_raw[byteIndex]=(byte)i;
                    byteIndex++;                
                }            
        }
        if("alpha".equals(this.tab.psychoPanel.byteGeneratorRange))
        {
                // 48-57 -> 10
                // 65-90 -> 26
                // 97-122 -> 26
                // 62 altogether
                this.bytes_raw = new byte[62];
                for(int i=48;i<58;i++)
                {
                    bytes_raw[byteIndex]=(byte)i;
                    byteIndex++;
                }
                for(int i=65;i<91;i++)
                {
                    bytes_raw[byteIndex]=(byte)i;
                    byteIndex++;                
                }
                for(int i=97;i<123;i++)
                {
                    bytes_raw[byteIndex]=(byte)i;
                    byteIndex++;                
                }
        }
        if("non-alpha-print".equals(this.tab.psychoPanel.byteGeneratorRange))
        {
                // 33-46 -> 14
                // 58-64 -> 7
                // 91-96 -> 6
                // 123-126 -> 4
                // 31 all together
                this.bytes_raw = new byte[31];
                for(int i=33;i<47;i++)
                {
                    bytes_raw[byteIndex]=(byte)i;
                    byteIndex++;
                }
                for(int i=58;i<65;i++)
                {
                    bytes_raw[byteIndex]=(byte)i;
                    byteIndex++;                
                }
                for(int i=91;i<97;i++)
                {
                    bytes_raw[byteIndex]=(byte)i;
                    byteIndex++;                
                }
                for(int i=123;i<127;i++)
                {
                    bytes_raw[byteIndex]=(byte)i;
                    byteIndex++;                
                }                
        }
        if("non-alpha-non-print-low".equals(this.tab.psychoPanel.byteGeneratorRange))
        {            
                // 0-31 -> 32
                // 127-128 -> 1
                // alltogether 33
                this.bytes_raw = new byte[33];
                for(int i=0;i<32;i++)
                {
                    bytes_raw[byteIndex]=(byte)i;
                    byteIndex++;                
                }
                for(int i=127;i<128;i++)
                {
                    bytes_raw[byteIndex]=(byte)i;
                    byteIndex++;                
                }
        }        
        if("non-alpha-non-print".equals(this.tab.psychoPanel.byteGeneratorRange))
        {            
                // 0-31 -> 32
                // 127-255 -> 129
                // alltogether 161
                this.bytes_raw = new byte[161];
                for(int i=0;i<32;i++)
                {
                    bytes_raw[byteIndex]=(byte)i;
                    byteIndex++;                
                }
                for(int i=127;i<256;i++)
                {
                    bytes_raw[byteIndex]=(byte)i;
                    byteIndex++;                
                }
        }
        if("all".equals(this.tab.psychoPanel.byteGeneratorRange))
        {            
                // 0-255 -> 256
                this.bytes_raw = new byte[256];
                for(int i=0;i<256;i++)
                {
                    bytes_raw[byteIndex]=(byte)i;
                    byteIndex++;                
                }
        }            
        // now, check what the encodings are - and fill the bytes_out string array accordingly
        // we simply use the global psychoPATH encoding options
       // ListModel encodeModel = this.tab.psychoPanel.encodeList.getModel();
        for(int j=0;j<this.bytes_raw.length;j++)
        {
                byte[] t = new byte[1];
                t[0]=this.bytes_raw[j];
                String convertedBytes = callbacks.getHelpers().bytesToString(t);
         //       for(int i=0;i<encodeModel.getSize();i++)
         //       {
         //          if("None".equals(encodeModel.getElementAt(i).toString()))
         //         { 
                       this.bytes_out.add(convertedBytes); // add the conversion here
         //          }
        //        }
        }
    }
    @Override
    public boolean hasMorePayloads() 
    {
            this.tab.psychoPanel.logOutput("Being asked for more payloads (payloadIndex:"+this.payloadIndex+" with size: "+this.bytes_out.size());
            return this.payloadIndex < this.bytes_out.size();
    }

    @Override
    public byte[] getNextPayload(byte[] baseValue) 
    {                
        byte[] payload;
        payload = this.callbacks.getHelpers().stringToBytes(this.bytes_out.get(this.payloadIndex).toString());        
        this.payloadIndex++; // increase the index
        return payload;
    }
    @Override
    public void reset() 
    {
        this.payloadIndex = 0;
    }       
}