/*## (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.io.*;
import java.util.*;
import java.security.*;

public class SCAsm {
 public static final int DSTART      = 0x4000;
 public static final int IOSTART     = 0x5e00;

 public static final int OUTPUT_TEXT     = 0x01;
 public static final int OUTPUT_IMAGE    = 0x02;
 public static final int OUTPUT_FILE     = 0x03;

 static Hashtable labels=new Hashtable();
 static Hashtable comments=new Hashtable();
 static Hashtable syms=new Hashtable();

 static Vector initdata=new Vector();

 static int codebase=0;
 static int database=DSTART;

 static int defidx=0;

 static int prevdefidx=-1;

 static void add_label(String name,int off) {
  if (labels.get(name)==null) {
   labels.put(name,new Integer(off));
  } else error("duplicate label: "+name);
 }

 static int get_label(String name) {
  Integer ival=(Integer)labels.get(name);

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

 public static void add_comment(int pos,String s) {
  if (comments.get(new Integer(pos))==null) {
   comments.put(new Integer(pos),s);
  }
 }

 public static String get_comment(int pos) {
  String s=(String)comments.get(new Integer(pos));

  return s;
 }

 public static void add_sym(String name,int val) {
  syms.put(name,new Integer(val));
 }

 public static int lookup_sym(String name) {
  Integer ival=(Integer)syms.get(name);

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

 public static void add_data(int addr,int val) {
  initdata.add(new Integer(addr));
  initdata.add(new Integer(val)); 
 }

 public static void init_syms() {
  try {
   syms=new Hashtable();

   add_sym("DMA_CONFIG",0x4004);
   add_sym("DMA_SRC",0x4008);
   add_sym("DMA_DST",0x400c);
   add_sym("DMA_SIZE",0x4010);
   add_sym("TK_CONFIG",0x4018);
   add_sym("TK_CMD",0x401c);
   add_sym("TK_DATA0",0x4020);
   add_sym("TK_DATA1",0x4024);
   add_sym("TK_DATA2",0x4028);
   add_sym("TK_DATA3",0x402c);
   add_sym("CUSTOMER_MODE",0x4030);

   add_sym("PUBLIC_ID",0x5e2c);
  } catch(Throwable t) {}
 }

 public static String hex_value(int val,int max) {
  String s="";
  for(int i=0;i<max;i++) {
   s+="0";
  }

  s+=Integer.toHexString(val);

  return s.substring(s.length()-max);
 }

 private static String do_pad(int cnt) {
  String s="";

  for(int i=0;i<cnt;i++) {
   s+=" ";
  }

  return s;
 }

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

 public static void handle_special(Token[] tokens,Token[] intokens) {
  String ins=tokens[0].val;

  if (ins.equals(".code")) {
   Token t1=tokens[1];
   Token t2=intokens[1];

   int val=Opcode.imm_val(t1.val,t2.val);
   Opcode.Encoder.set_codebase(val);
   System.out.println(".code 0x"+hex_value(val,4));
  } else

  if (ins.equals(".data")) {
   Token t1=tokens[1];
   Token t2=intokens[1];

   int val=Opcode.imm_val(t1.val,t2.val);
   Opcode.Encoder.set_database(val);
   System.out.println(".data 0x"+hex_value(val,4));
  } else

  if (ins.equals(".equ")) {
   Token t0=intokens[1];

   Token t1=tokens[2];
   Token t2=intokens[2];

   int val=Opcode.imm_val(t1.val,t2.val);
 
   String sym=t0.val;

   add_sym(sym,val);

   System.out.println(".equ "+sym+"=0x"+hex_value(val,4));   
  } else

  if (ins.equals(".def")) {
   Token t0=intokens[1];
   String sym=t0.val;

   int val=Opcode.Encoder.database+4*defidx;
   add_sym(sym,val);

   prevdefidx=defidx;   
   defidx++;
  } else

  if (ins.equals(".word")) {
   if (prevdefidx==(defidx-1)) {
    defidx--;
   }

   Token t1=tokens[1];
   Token t2=intokens[1];

   int val=Opcode.imm_val(t1.val,t2.val);
   int addr=Opcode.Encoder.database+4*defidx;

   add_data(addr,val);

   defidx++;
  }
 }

 public static int[] process(String lines[],boolean firstpass) {
  int off=0;
  int linenum=1;

  Vector v=new Vector();

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

   if (line.equals("")||line.startsWith(";")) {
    linenum++;
    continue;
   }

   Token tokens[]=Token.tokenize(line);

   if (tokens==null) error("parsing line "+linenum);

   if ((tokens.length==1)&&(tokens[0].type==Token.TOK_LABEL)) {
    if (firstpass) {
     String label=tokens[0].val;
     label=label.substring(0,label.length()-1);

     add_label(label,off);
    }
   } else
    
   if (tokens.length>0) {
    Opcode.Encoder encoder=Opcode.Encoder.find(tokens,linenum);

    if (encoder==null) error("invalid instruction format in line "+linenum);

    if (!firstpass) {
     String ins=tokens[0].val;

     if (special_ins(ins)) {
      handle_special(encoder.tokens,tokens);
     } else {
      int opcode=encoder.build(tokens);
      v.add(new Integer(opcode));
      add_comment(off,line);
      off++;
     }
    } else {
     String ins=tokens[0].val;
     if (!special_ins(ins)) off++;
    }
   } else error("parsing line "+linenum);

   linenum++;
  }

  int opcodes[]=new int[v.size()];

  for(int i=0;i<opcodes.length;i++) {
   opcodes[i]=((Integer)v.elementAt(i)).intValue();
  }

  return opcodes;
 }

 public static int[] assemble(String filename) {
  int opcodes[]=null;

  try {
   String lines[]=load_file(filename);

   if (lines==null) error("reading file "+filename);

   labels=new Hashtable();
   Opcode.Encoder.set_codebase(codebase);
   Opcode.Encoder.set_database(database);

   process(lines,true);
   opcodes=process(lines,false);
  } catch(Throwable t) {}

  return opcodes;
 }

 private static void dump_code(String filename,boolean image) {
  String line=null;

  System.out.println("- assembling");
  int opcodes[]=assemble(filename);

  System.out.println("[CODE codebase=0x"+hex_value(Opcode.Encoder.codebase,4)+", database=0x"+hex_value(Opcode.Encoder.database,4)+"]");

  if (image) {
   line=do_pad(2)+"0x"+hex_value(opcodes.length,8)+","+do_pad(4)+"// opcodes length";
   System.out.println(line);
   line=do_pad(2)+"//.code";
   System.out.println(line);
  }

  for(int i=0;i<opcodes.length;i++) {
   String comma="";
   if ((i<(opcodes.length-1))||(initdata.size()>0)) comma+=",";
    else comma+=" ";

   line=do_pad(2)+"0x"+hex_value(opcodes[i],8)+comma+do_pad(4)+"// 0x"+hex_value(Opcode.Encoder.codebase+i,4)+do_pad(2)+get_comment(i);
   System.out.println(line);
  }

  if (image) {
   line=do_pad(2)+"//.initdata";
   System.out.println(line);

   line=do_pad(2)+"0x"+hex_value(initdata.size(),8)+","+do_pad(4)+"// initdata length";
   System.out.println(line);

   for(int i=0;i<initdata.size();i+=2) {
    int addr=((Integer)initdata.elementAt(i)).intValue();
    int val=((Integer)initdata.elementAt(i+1)).intValue();

    line=do_pad(2)+"0x"+hex_value(addr,8)+","+do_pad(4)+"// addr";
    System.out.println(line);

    String comma="";
    if ((i+2)<initdata.size()) comma+=",";
     else comma+=" ";

    line=do_pad(2)+"0x"+hex_value(val,8)+comma+do_pad(4)+"// val";
    System.out.println(line);
   }
  }
 }

 private static void save_code(String filename,String outfile) {
  System.out.println("- assembling");
  int opcodes[]=assemble(filename);

  System.out.println("[CODE codebase=0x"+hex_value(Opcode.Encoder.codebase,4)+", database=0x"+hex_value(Opcode.Encoder.database,4)+"]");

  int idata_size=initdata.size();

  int idata[]=new int[idata_size];

  for(int i=0;i<initdata.size();i+=2) {
   int addr=((Integer)initdata.elementAt(i)).intValue();
   int val=((Integer)initdata.elementAt(i+1)).intValue();

   idata[i]=addr;   
   idata[i+1]=val; 
  }

  save(outfile,opcodes,idata);
 }

 private static String[] load_file(String name) {
  String res[]=null;

  try {
   System.out.println("- loading "+name);

   Vector v=new Vector();

   FileInputStream fis=new FileInputStream(name);

   BufferedReader b=new BufferedReader(new InputStreamReader(fis,"utf-8"));

   String line=b.readLine();

   while(line!=null) {
    line=line.trim();
    v.add(line);
    line=b.readLine();
   }

   b.close();
   fis.close();

   res=new String[v.size()];

   for(int i=0;i<res.length;i++) {
    res[i]=(String)v.elementAt(i);    
   }
  } catch(Throwable e) {
   e.printStackTrace();
  }

  return res;
 }

 private static void save_int(DataOutputStream dos,int val) throws Throwable {
  dos.write(val&0xff);
  dos.write((val>>8)&0xff);
  dos.write((val>>16)&0xff);
  dos.write((val>>24)&0xff);
 }

 private static void save(String name,int cbuf[],int dbuf[]) {
  FileOutputStream fos=null;

  System.out.println("- saving "+name);

  try {
   fos=new FileOutputStream(name,false);
  } catch(Throwable t) {}

  if (fos==null) error("cannot create file");

  try {
   DataOutputStream dos=new DataOutputStream(fos);

   save_int(dos,cbuf.length);

   for(int i=0;i<cbuf.length;i++) {
    int val=cbuf[i];
    save_int(dos,val);
   }

   save_int(dos,dbuf.length);

   for(int i=0;i<dbuf.length;i++) {
    int val=dbuf[i];
    save_int(dos,val);
   }

   dos.close();
   fos.close();
  } catch(Throwable t) {error("cannot write to file");}
 }

 public static void usage() {
  System.out.println("usage: SCAsm -f src_name [-c code_base][-d data_base][-t][-i][-o dump_file]");
  System.exit(0);
 }

 public static void error(String s) {
  System.out.println("error: "+s);
  System.exit(0);
 }
 
 public static void check_idx(String args[],int idx) {
  if (idx>=args.length) usage();
 }

 public static void main(String args[]) throws Throwable {
  try {
   System.out.println("/*## (c) SECURITY EXPLORATIONS    2018 poland                                #*/");
   System.out.println("/*##     http://www.security-explorations.com                                #*/");
   System.out.println("SlimCore assembler");

   int cmd=0;
   int idx=0;

   String src_filename=null;
   String out_filename=null;

   if (args.length==0) usage();

   while(idx<args.length) {
    String arg=args[idx++];

    if (arg.equals("-f")) {
     check_idx(args,idx);
     src_filename=args[idx++];
    } else

    if (arg.equals("-c")) {
     check_idx(args,idx);
     codebase=Integer.parseInt(args[idx++],16);
    } else

    if (arg.equals("-d")) {
     check_idx(args,idx);
     database=Integer.parseInt(args[idx++],16);
    } else

    if (arg.equals("-t")) {
     cmd=OUTPUT_TEXT;
    } else

    if (arg.equals("-i")) {
     cmd=OUTPUT_IMAGE;
    } else

    if (arg.equals("-o")) {
     check_idx(args,idx);
     out_filename=args[idx++];
     cmd=OUTPUT_FILE;
    } else usage();
   }

   if (cmd==0)  error("no command");

   if (src_filename==null) error("no source file");

   init_syms();

   switch(cmd) {
    case OUTPUT_TEXT:
     dump_code(src_filename,false);
     break;
    case OUTPUT_IMAGE:
     dump_code(src_filename,true);
     break;
    case OUTPUT_FILE:
     save_code(src_filename,out_filename);
     break;
   }
  } catch(Throwable e) {
   e.printStackTrace();
  } 
 }
}
