/*## (c) SECURITY EXPLORATIONS    2018 poland                                #*/
/*##     http://www.security-explorations.com                                #*/

/* RESEARCH MATERIAL:	SRP-2018-02                                           */
/* SlimCORE assembler                                                         */

/* THIS SOFTWARE IS PROTECTED BY DOMESTIC AND INTERNATIONAL COPYRIGHT LAWS    */
/* UNAUTHORISED COPYING OF THIS SOFTWARE IN EITHER SOURCE OR BINARY FORM IS   */
/* EXPRESSLY FORBIDDEN. ANY USE, INCLUDING THE REPRODUCTION, MODIFICATION,    */
/* DISTRIBUTION, TRANSMISSION, RE-PUBLICATION, STORAGE OR DISPLAY OF ANY      */
/* PART OF THE SOFTWARE, FOR COMMERCIAL OR ANY OTHER PURPOSES REQUIRES A      */
/* VALID LICENSE FROM THE COPYRIGHT HOLDER.                                   */

/* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS    */
/* OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,*/
/* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL    */
/* SECURITY EXPLORATIONS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, */
/* WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF  */
/* OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE     */
/* SOFTWARE.                                                                  */

import java.lang.*;
import java.util.*;

public class Opcode {
 static String[] map[]={
  //00 opcodes
  {"mov",  "0000 reg1 0000 reg2 0011 1100", "reg1 , reg2"},
  {"swap", "0000 reg1 reg2 0000 1100 0000", "reg1 , reg2"},

  //01 opcodes
  {"shl",  "0001 reg1 reg2 0000 000 imm5",  "reg1 , reg2 , # imm5"},
  {"shr",  "0001 reg1 reg2 0000 001 imm5",  "reg1 , reg2 , # imm5"},

  //02 opcodes
  {"add",  "0010 reg1 reg2 reg3 imm8", "reg1 , reg2 , reg3 , # imm8"},

  //03 opcodes
  {"sub",  "0011 reg1 reg2 reg3 imm8", "reg1 , reg2 , reg3 , # imm8"},

  //04 opcodes
  {"and", "0100 reg1 reg2 reg3 00000000", "reg1 , reg2 , reg3"},
  {"and", "0100 reg1 reg2 0000 imm8",     "reg1 , reg2 , # imm8"},
  {"tst", "0100 0000 reg  0000 imm8",     "reg , # imm8"},
  {"tst", "0100 0000 reg1 reg2 00000000", "reg1 , reg2"},

  //05 opcodes
  {"or",  "0101 reg1 reg2 reg3 00000000", "reg1 , reg2 , reg3"},
  {"or",  "0101 reg1 reg2 0000 imm8",     "reg1 , reg2 , # imm8"},
  {"tst", "0101 0000 0000 reg 00000000",  "reg , reg"},

  //06 opcodes
  {"xor", "0110 reg1 reg2 reg3 00000000", "reg1 , reg2 , reg3"},
  {"xor", "0110 reg1 reg2 0000 imm8",     "reg1 , reg2 , # imm8"},

  //07 opcodes
  {"and",    "0111 reg1 reg2 11 bitnum 00000", "reg1 , reg2 , bitnum"},
  {"mov",    "0111 reg1 reg2 11 bitnum shift", "reg1 , ( reg2 >> shift ) & bitnum"},
  {"mov",    "0111 reg1 reg2 00 bitnum shift", "reg1 , ( reg2 & bitnum ) << shift"},
  {"mov",    "0111 reg1 reg2 01 bitnum 00000", "reg1 , reg2 & bitnum"},
  {"mov",    "0111 reg1 reg2 00 bitnum 00000", "reg1 , reg2 & bitnum"},
  {"mov",    "0111 reg1 reg2 01 00001  shift", "reg1 , reg2 &0x01 << shift"},
  {"movzx",  "0111 reg1 reg2 01 01000 00000",  "reg1 , reg2 &0xff"},
  {"movhi",  "0111 reg1 reg2 00 bitnum 10000", "reg1 , ( reg2 & bitnum ) <<16"},
  {"movhi",  "0111 reg1 reg2 00 10000 10000",  "reg1 , reg2 <<16"},
  {"bitset", "0111 reg1 reg2 00 00001 shift",  "reg1 , reg2 &0x01 << shift"},
  {"bitclr", "0111 reg1 0000 00 00001 shift",  "reg1 , #1 << shift"},
  {"bitval", "0111 reg1 reg2 11 00001 shift",  "reg1 , reg2 , #1 << shift"},
  {"bittst", "0111 0000 reg 11 00001 shift",   "reg , #1 << shift"},

  //08 opcodes
  {"jmp",   "1000 0100 0000 reg 00000000", "reg"},
  {"j",     "1000 1100 0000 target",       "target"},
  {"jz",    "1000 1000 0001 target",       "target"},
  {"jne",   "1000 1100 0001 target",       "target"},
  {"js",    "1000 1000 0010 target",       "target"},
  {"jns",   "1000 1100 0010 target",       "target"},
  {"jxx1",  "1000 1111 0010 target",       "target"},
  {"wait1", "1000 1110 0001 target",       "target"},
  {"wait2", "1000 1101 1000 target",       "target"},

  //09 opcodes
  {"je",    "1001 1000 0001 target", "target"},
  {"jb",    "1001 1000 0100 target", "target"},
  {"jae",   "1001 1100 0100 target", "target"},
  {"jbe",   "1001 1000 0111 target", "target"},
  {"jne",   "1001 1100 0001 target", "target"},
  {"jns",   "1001 1100 0010 target", "target"},
  {"js",    "1001 1000 0010 target", "target"},
  {"jzs",   "1001 1000 0011 target", "target"},
  {"jxx2",  "1001 1011 0100 target", "target"},
  {"jxx3",  "1001 1010 1000 target", "target"},
  {"wait3", "1001 1110 0001 target", "target"},
  {"wait4", "1001 1101 1000 target", "target"},

  //0a opcodes
  {"ld", "1010 reg1 0000 reg2 mimm8", "reg1 , [ reg2 + mimm8 ]"},
  {"ld", "1010 reg1 0001 reg2 imm8",  "reg1 , [ reg2 ] , reg2 +=# imm8"},

  //0b opcodes
  {"st", "1011 0000 reg1 reg2 mimm8", "reg1 , [ reg2 + mimm8 ]"},
  {"st", "1011 0001 reg1 reg2 imm8",  "reg1 , [ reg2 ] , reg2 +=# imm8"},

  //0c opcodes
  {"cmp", "1100 0000 reg 0000 imm8", "reg , # imm8"},

  //0d opcodes
  {"jmp",             "1101 0000 target_hi 0001 target_lo", "target"},
  {"bitsrch_topmost", "1101 reg1 reg2 0000 0100 0000",      "reg1 , reg2"},
  {"sync1",           "1101 0000 0000 0000 1000 0000",      ""},
  {"sync2",           "1101 0000 0000 0000 1001 0000",      ""},
  {"rpt",             "1101 0000 0000 0000 0000 n",         "n"},
  {"rpt_16",          "1101 0000 0001 0000 0000 0000",      ""},

  //0e opcodes
  {"mov", "1110 reg imm16", "reg , # imm16"},

  //0f opcodes
  {"copAES",      "1111 0101 0100 0000 0000 0000", ""},
  {"copTDES",     "1111 1010 0100 0000 0000 0000", ""},
  {"copAES_dma",  "1111 0100 0100 0000 0000 0000", ""},
  {"copTDES_dma", "1111 1000 0100 0000 0000 0000", ""},

  //generic opcode
  {"opcode",      "imm24", "imm24"},

  //special opcodes
  {".code",    "0000 0000 0000 0000 0000 0000", "imm16"},
  {".data",    "0000 0000 0000 0000 0000 0000", "imm16"},
  {".equ",     "0000 0000 0000 0000 0000 0000", "sym imm16"},
  {".def",     "0000 0000 0000 0000 0000 0000", "sym"},
  {".word",    "0000 0000 0000 0000 0000 0000", "imm32"}
 };

 static {
  components=new Hashtable();
  add_component("reg",4);
  add_component("reg1",4);
  add_component("reg2",4);
  add_component("reg3",4);
  add_component("imm5",5);
  add_component("imm8",8);
  add_component("mimm8",8);
  add_component("imm16",16);
  add_component("shift",5);
  add_component("bitnum",5);
  add_component("target",12);
  add_component("target_hi",8);
  add_component("target_lo",4);
  add_component("n",4);

  add_component("imm24",24);

  add_component("imm32",32);

  build_encoders();
  build_helper();
 }

 static Hashtable components;

 static Encoder encoders[];
 static Hashtable helper;

 static boolean pureval;

 static void error(String err) {
  System.out.println("error: "+err);
  System.exit(1);
 }

 static void usage() {
  System.out.println("usage: SCAsm filename base");
  System.exit(1);
 }

 public static int parse_hex(String s) {
  int val=0;

  for(int i=0;i<s.length();i++) {
   char ch=s.charAt(i);
   int d;

   if ((ch>='0')&&(ch<='9')) {
    d=ch-'0';
   } else

   if ((ch>='a')&&(ch<='f')) {
    d=10+ch-'a';
   } else

   if ((ch>='A')&&(ch<='F')) {
    d=10+ch-'A';
   } else return -1;

   val=(val<<4)|(d&0x0f);
  }

  return val;
 }

 static void build_encoders() {
  int size=map.length;

  encoders=new Encoder[size];

  for(int i=0;i<size;i++) {
   encoders[i]=new Encoder(map[i]);
  }
 }

 static void build_helper() {
  helper=new Hashtable();

  for(int i=0;i<encoders.length;i++) {
   Encoder encoder=encoders[i];

   EncoderItem ei=(EncoderItem)helper.get(encoder.name);

   if (ei==null) {
    ei=new EncoderItem(encoder);
    helper.put(encoder.name,ei);
   } else {
    ei.add(encoder);
   }
  }
 }

 static void add_component(String name,int size) {
  components.put(name,new Integer(size));
 }

 static int component_size(String name) {
  Integer ival=(Integer)components.get(name);

  if (ival!=null) return ival.intValue();
   else return 0;
 }

 public static int reg_val(String name,String tok_val) {
  int reg=-1;

  if (!tok_val.startsWith("r")) return reg;

  try {
   String rname=tok_val.substring(1);
   reg=Integer.parseInt(rname);
   if ((reg<0)||(reg>16)) reg=-1;
  } catch(Throwable t) {}

  return reg;
 }

 public static int target_val(String name,String tok_val) {
  int target=-1;

  try {
   if (tok_val.startsWith("0x")) {
    target=Integer.parseInt(tok_val.substring(2),16);
   } else
   if (tok_val.startsWith("0")) {
    target=Integer.parseInt(tok_val,16);
   } else {
    int ch=(int)tok_val.getBytes()[0];

    if (Token.is_digit(ch)) {
     target=Integer.parseInt(tok_val);
    } else {
     target=SCAsm.get_label(tok_val);  

     if (target<0) error("unknown label: "+tok_val);

     target+=Encoder.codebase;
    }
   }
  } catch(Throwable t) {}

  if (name.equals("target")) {
   target&=0xfff;
  } else

  if (name.equals("target_hi")) {
   target=(target>>4)&0xff;
  } else

  if (name.equals("target_lo")) {
   target&=0xf;
  }

  return target;
 }

 public static int imm_val(String name,String tok_val) {
  int imm=-1;

  int size=component_size(name);

  int max_val=(1<<size)-1;

  try {
   if (tok_val.startsWith("0x")) {
    //imm=Integer.parseInt(tok_val.substring(2),16);
    imm=parse_hex(tok_val.substring(2));
   } else
   if (tok_val.startsWith("0")) {
    imm=Integer.parseInt(tok_val,16);
   } else {
    int ch=(int)tok_val.getBytes()[0];

    if (Token.is_digit(ch)) {
     imm=Integer.parseInt(tok_val);
    } else {
     int sval=SCAsm.lookup_sym(tok_val);
 
     if (sval==-1) {
      sval=SCAsm.get_label(tok_val);  

      if (sval<0) error("unknown symbol/label: "+tok_val);
      sval+=Encoder.codebase;
     }

     if (!pureval) {
      if ((sval>=SCAsm.DSTART)&&(sval<SCAsm.IOSTART)) {
       sval-=SCAsm.DSTART;
       sval>>>=2;
      } else

      if (sval>=SCAsm.IOSTART) {
       sval-=SCAsm.IOSTART;
       sval>>>=2;
       sval+=0x80;
      }
     }

     imm=sval;
    }
   }

//   if ((imm<-max_val)||(imm>max_val)) imm=-1;
  } catch(Throwable t) {
   t.printStackTrace();
  }

  return imm;
 }

 public static int bitnum_val(String name,String tok_val) {
  int mask=imm_val(name,tok_val);

  for(int i=0;i<Encoder.masks.length;i++) {
   if (Encoder.masks[i]==mask) return i+1;
  }

  return 0;
 }

 public static boolean is_reg_operand(String name) {
  if (name.startsWith("reg")) return true;
   else return false;
 }

 public static boolean is_imm_operand(String name) {
  String imm_operands[]={
   "imm5","imm8","mimm8","imm16","shift","bitnum","n","imm24","imm32"
  };

  for(int i=0;i<imm_operands.length;i++) {
   if (imm_operands[i].equals(name)) return true;
  }

  return false;
 }

 public static boolean is_target_operand(String name) {
  if (name.equals("target")) return true;
   else return false;
 }

 public static class ComponentBuilder {
  String name;
  int shift;
  int bitmask;

  public int build(String tokval) {
   int val=0;

   if (name.startsWith("reg")) {
    val=reg_val(name,tokval);
   } else

   if (name.startsWith("target")) {
    val=target_val(name,tokval);
   } else

   if (name.startsWith("bitnum")) {
    val=bitnum_val(name,tokval);
   } else {
    val=imm_val(name,tokval);
   }

   if (val<0) return val;

   val=(val&bitmask)<<shift;

   return val;
  }

  public ComponentBuilder(String name,int shift,int size) {
   this.name=name;
   this.shift=shift;
   this.bitmask=Encoder.mask(size-1);
  }
 }

 public static class Encoder {
  String name;
  String sformat;
  String sprint_tokens;
  String[] format;
  Token tokens[];
  int weight;
  int val;
  int mask;

  static int masks[];
  static int codebase;
  static int database;

  //component builders for this opcode format
  //a map of component name -> ComponentBuilder instance
  Hashtable builders;

  static {
   masks=new int[32];
   
   for(int i=0;i<masks.length;i++) {
    if (i==0) masks[i]=1;
     else masks[i]=masks[i-1]|(1<<i);
   }
  }

  public String[] str2tokens(String s) {
   s=s.trim();

   StringTokenizer st=new StringTokenizer(s," ", false);
   int cnt=st.countTokens();

   String tokens[]=new String[cnt];

   for(int i=0;i<cnt;i++) {
    tokens[i]=st.nextToken();
   }            

   return tokens;
  }

  public int weight(Token[] tokens) {
   int weight=0;

   for(int i=0;i<tokens.length;i++) {
    Token tok=tokens[i];

    if (tok.type==Token.TOK_NUMBER) weight++;
   }

   return weight;
  }

  public int weight(String[] format) {
   int weight=0;

   for(int i=0;i<format.length;i++) {
    String component=format[i];

    if (component.startsWith("0")||component.startsWith("1")) {
     weight+=component.length();
    }
   }

   return weight;
  }

  public static int mask(int bitnum) {
   return masks[bitnum];
  }

  public int mask(int bitnum,int len) {
   int mask=masks[bitnum];

   int nbitnum=bitnum-len-1;

   if (nbitnum>=0) {
    mask^=masks[nbitnum];
   }

   return mask;
  }

  public void calc_value_and_mask() {
   mask=0;
   val=0;

   int shift=24;

   for(int i=0;i<format.length;i++) {
    String component=format[i];

    int size=component_size(component);

    if (size>0) {
     shift-=size;
     val<<=size;
    } else {
     mask|=mask(shift-1,component.length()-1);
     shift-=component.length();
     val<<=component.length();
     val|=Integer.parseInt(component,2);
    }
   }   
  }

  public void setup_builders() {
   builders=new Hashtable();

   int shift=24;

   for(int i=0;i<format.length;i++) {
    String component=format[i];

    int size=component_size(component);

    if (size>0) {
     shift-=size;
//     System.out.println(" - "+component+" shift "+Integer.toHexString(shift)+" mask "+Integer.toHexString(mask(size-1)));
     builders.put(component,new ComponentBuilder(component,shift,size));
    } else {
     shift-=component.length();
    }
   }   
  }

  public boolean match(Token intokens[]) {
   if (tokens.length!=intokens.length) return false;

   for(int i=0;i<tokens.length;i++) {
    Token t1=tokens[i];
    Token t2=intokens[i];

    if (t1.type!=t2.type) {
     if (t1.type==Token.TOK_SYM) {
      boolean ok1=is_imm_operand(t1.val)&&(t2.type==Token.TOK_NUMBER);
      boolean ok2=is_target_operand(t1.val)&&(t2.type==Token.TOK_NUMBER);
      if ((!ok1)&&(!ok2)) return false;
     } else return false;
    } else

    if ((t1.type==Token.TOK_SYM)&&is_reg_operand(t1.val)) {
     int reg=reg_val(t1.val,t2.val);
     if (!(reg>=0)) return false;
    }
   }

   return true;
  }

  public int token_idx(String component) {
   for(int i=0;i<tokens.length;i++) {
    Token tok=tokens[i];

    if (tok.type==Token.TOK_SYM) {
     if (tok.val.equals(component)) return i;
    }
   }

   return -1;
  }

  public int build(Token intokens[]) {
   int opcode=val;

   for(int i=0;i<format.length;i++) {
    String component=format[i];

    int size=component_size(component);

    if (size>0) {
     int idx=token_idx(component);

     if (idx<0) {
      if (component.startsWith("target_")) {
       idx=token_idx("target");
      }
     }

     if (idx>0) {
      ComponentBuilder cb=(ComponentBuilder)builders.get(component);
      String tok_val=intokens[idx].val;

      if (idx>0) {
       int prevtok_type=intokens[idx-1].type;
       if (prevtok_type==Token.TOK_HASH) pureval=true;
        else pureval=false;
      }

      opcode|=cb.build(tok_val);
     }
    }
   }

   return opcode;
  }

  public Encoder(String name,String sformat,String sprint_tokens) {
   this.name=name;
   this.sformat=sformat;
   this.sprint_tokens=sprint_tokens;
   this.format=str2tokens(sformat);
   this.tokens=Token.tokenize(name+" "+sprint_tokens);

   weight=weight(format);
   calc_value_and_mask();
   setup_builders();
  }

  public Encoder(String[] opcode_desc) {
   this(opcode_desc[0],opcode_desc[1],opcode_desc[2]);
  }

  public static void set_codebase(int cbase) {
   codebase=cbase;
  }

  public static void set_database(int dbase) {
   database=dbase;
  }

  public static Encoder find(Token tokens[],int linenum) {
   String ins=tokens[0].val;

   EncoderItem ei=(EncoderItem)helper.get(ins);

   if (ei==null) error("unknown instruction at line "+linenum);

   while(ei!=null) {
    Encoder encoder=ei.encoder;

    boolean found=encoder.match(tokens);

    if (found) {
     if (ins.equals("jmp")) {
      String op=encoder.tokens[1].val;

      if (is_reg_operand(op)&&(reg_val(op,tokens[1].val)>=0)) return encoder;
       else
      if (is_target_operand(op)) return encoder;
       else ei=ei.next;
     } else return encoder;
    } else ei=ei.next;
   }

   return null;
  }
 }

 public static class EncoderItem {
  Encoder encoder;
  EncoderItem next;

  public EncoderItem(Encoder encoder) {
   this.encoder=encoder;
  }

  public EncoderItem(Encoder encoder,EncoderItem next) {
   this.encoder=encoder;
   this.next=next;
  }

  public void add(Encoder d) {
   if (encoder.weight>d.weight) {
    if (next!=null) {
     next.add(d);
    } else {
     next=new EncoderItem(d);
    }
   } else {
    next=new EncoderItem(encoder,next);
    encoder=d;
   }
  }
 }

}
