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

/* 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.*;

public class GenStubs {
 public static class MethodInfo {
  int access;
  String name;
  String desc;

  public MethodInfo(int access,String name,String desc) {
   this.access=access;
   this.name=name;
   this.desc=desc;
  }
 }

 public static class FieldInfo {
  int access;
  String name;
  String type;
  int val;

  public FieldInfo(int access,String name,String type) {
   this.access=access;
   this.name=name;
   this.type=type;
  }

  public FieldInfo(int access,String name,String type,int val) {
   this.access=access;
   this.name=name;
   this.type=type;
   this.val=val;
  }
 }

 public static class ClassInfo {
  int access;
  String name;
  String supername;
  Vector interfaces;
  Vector methods;
  Vector fields;

  public ClassInfo(int access,String name,String supername) {
   this.access=access;
   this.name=name;
   if (supername.equals("null")) supername=null;
   this.supername=supername;
   interfaces=new Vector();
   methods=new Vector();
   fields=new Vector();
  }
 }

 public static final String STUBS_DIR   = "..\\stubs";

 public static final int ACC_PUBLIC     = 0x0001;
 public static final int ACC_PROTECTED  = 0x0004;
 public static final int ACC_STATIC     = 0x0008;
 public static final int ACC_FINAL      = 0x0010;
 public static final int ACC_TRANSIENT  = 0x0080;
 public static final int ACC_NATIVE     = 0x0100;
 public static final int ACC_INTERFACE  = 0x0200;
 public static final int ACC_ABSTRACT   = 0x0400;

 public static final int CLASS_MAGIC    = 0xCAFEBABE;
 public static final int CLASS_MAJOR    = 0x2e;
 public static final int CLASS_MINOR    = 0x00;

 public static final int CONSTANT_Utf8      = 1;
 public static final int CONSTANT_Integer   = 3;
 public static final int CONSTANT_Class     = 7;

 public static final int ST_INIT            = 0;
 public static final int ST_CLASS_BLOCK     = 1;
 public static final int ST_INTERFACE_BLOCK = 2;
 public static final int ST_METHOD_BLOCK    = 3;
 public static final int ST_FIELD_BLOCK     = 4;

 public static void usage() {
  System.out.println("usage: GenStubs classdesc_file");
  System.exit(0);
 }

 public static void error(String s) {
  System.out.println("error: "+s);
  System.exit(0);
 }

 public static void error_with_line(String s,int num) {
  error(s+" [line "+num+"]");
 }

 public static 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 static ClassInfo get_class(String tokens[],int lineno) {
  if (tokens.length<5) error_with_line("invalid classdesc file",lineno);

  String type=tokens[tokens.length-5];
  String classname=tokens[tokens.length-4];
  String etoken=tokens[tokens.length-3];
  String supername=tokens[tokens.length-2];
  String ebracket=tokens[tokens.length-1];

  if (!etoken.equals("extends")) error_with_line("expected extends after classname",lineno);
  if (!ebracket.equals("{")) error_with_line("expected {",lineno);

  int access=0;

  if (type.equals("interface")) {
   access|=ACC_INTERFACE;
  } else
  if (!type.equals("class")) error_with_line("expected class or interface",lineno);

  for(int i=0;i<tokens.length-5;i++) {
   String access_token=tokens[i];

   if (access_token.equals("public")) {
    access|=ACC_PUBLIC;
   } else

   if (access_token.equals("final")) {
    access|=ACC_FINAL;
   } else

   if (access_token.equals("abstract")) {
    access|=ACC_ABSTRACT;
   } else error_with_line("invalid access flags",lineno);
  }

  return new ClassInfo(access,classname,supername);
 }

 public static int parse_class_block(String tokens[],int lineno) {
  if (tokens.length==1) {
   if (tokens[0].equals("}")) return ST_INIT;
  }

  if (tokens.length<2) error_with_line("invalid classdesc file",lineno);

  String type=tokens[tokens.length-2];
  String ebracket=tokens[tokens.length-1];

  if (!ebracket.equals("{")) error_with_line("expected {",lineno);

  if (type.equals("interfaces")) return ST_INTERFACE_BLOCK;
   else
  if (type.equals("methods")) return ST_METHOD_BLOCK;
   else
  if (type.equals("fields")) return ST_FIELD_BLOCK;
   else error_with_line("expected interfaces or methods",lineno);

  return -1;
 }

 public static int parse_interface(ClassInfo clazz,String tokens[],int lineno) {
  if (tokens.length!=1) error_with_line("invalid classdesc file",lineno);

  String name=tokens[tokens.length-1];

  if (name.equals("}")) return ST_CLASS_BLOCK;
  
  clazz.interfaces.add(name);

  return ST_INTERFACE_BLOCK;
 }

 public static int parse_method(ClassInfo clazz,String tokens[],int lineno) {
  if (tokens.length==1) {
   if (tokens[0].equals("}")) return ST_CLASS_BLOCK;
  }

  if (tokens.length<2) error_with_line("invalid classdesc file",lineno);

  String nametype=tokens[tokens.length-1];

  int access=0;

  for(int i=0;i<tokens.length-1;i++) {
   String access_token=tokens[i];

   if (access_token.equals("public")) {
    access|=ACC_PUBLIC;
   } else

   if (access_token.equals("protected")) {
    access|=ACC_PROTECTED;
   } else

   if (access_token.equals("static")) {
    access|=ACC_STATIC;
   } else

   if (access_token.equals("final")) {
    access|=ACC_FINAL;
   } else

   if (access_token.equals("native")) {
    access|=ACC_NATIVE;
   } else

   if (access_token.equals("abstract")) {
    access|=ACC_ABSTRACT;
   } else error_with_line("invalid access flags",lineno);
  }

  int idx=nametype.indexOf("(");
  if (idx<0) error_with_line("invalid method desc",lineno);

  String name=nametype.substring(0,idx);
  String desc=nametype.substring(idx,nametype.length());

  clazz.methods.add(new MethodInfo(access,name,desc));

  return ST_METHOD_BLOCK;
 }

 public static int parse_field(ClassInfo clazz,String tokens[],int lineno) {
  if (tokens.length==1) {
   if (tokens[0].equals("}")) return ST_CLASS_BLOCK;
  }

  if (tokens.length<3) error_with_line("invalid classdesc file",lineno);

  String equal_token=tokens[tokens.length-2];

  if (equal_token.equals("=")) {
   if (tokens.length<5) error_with_line("invalid classdesc file",lineno);

   String type=tokens[tokens.length-4];
   String name=tokens[tokens.length-3];
   String sval=tokens[tokens.length-1];

   int access=0;

   for(int i=0;i<tokens.length-4;i++) {
    String access_token=tokens[i];

    if (access_token.equals("public")) {
     access|=ACC_PUBLIC;
    } else

    if (access_token.equals("protected")) {
     access|=ACC_PROTECTED;
    } else

    if (access_token.equals("static")) {
     access|=ACC_STATIC;
    } else

    if (access_token.equals("final")) {
     access|=ACC_FINAL;
    } else 

    if (access_token.equals("transient")) {
     access|=ACC_TRANSIENT;
    } else error_with_line("invalid access flags",lineno);
   }

   int ival=Integer.parseInt(sval);
   clazz.fields.add(new FieldInfo(access,name,type,ival));
  } else {
   String type=tokens[tokens.length-2];
   String name=tokens[tokens.length-1];

   int access=0;

   for(int i=0;i<tokens.length-2;i++) {
    String access_token=tokens[i];

    if (access_token.equals("public")) {
     access|=ACC_PUBLIC;
    } else

    if (access_token.equals("protected")) {
     access|=ACC_PROTECTED;
    } else

    if (access_token.equals("static")) {
     access|=ACC_STATIC;
    } else

    if (access_token.equals("final")) {
     access|=ACC_FINAL;
    } else 

    if (access_token.equals("transient")) {
     access|=ACC_TRANSIENT;
    } else error_with_line("invalid access flags",lineno);
   }

   clazz.fields.add(new FieldInfo(access,name,type));
  }

  return ST_FIELD_BLOCK;
 }

 public static Vector load_classdesc(String name) {
  Vector classes=new Vector();
  int state=ST_INIT;
  int next_state=state;

  try {
   FileInputStream fis=new FileInputStream(name);
   BufferedReader b=new BufferedReader(new InputStreamReader(fis,"utf-8"));

   int lineno=1;

   String line=b.readLine();

   ClassInfo cur_class=null;

   while(line!=null) {
    String tokens[]=str2tokens(line);

    if (tokens.length!=0) {
     switch(state) {
      case ST_INIT:
       cur_class=get_class(tokens,lineno);
       next_state=ST_CLASS_BLOCK;
       break;
      case ST_CLASS_BLOCK:
       next_state=parse_class_block(tokens,lineno);
       break;
      case ST_INTERFACE_BLOCK:
       next_state=parse_interface(cur_class,tokens,lineno);
       break;
      case ST_METHOD_BLOCK:
       next_state=parse_method(cur_class,tokens,lineno);
       break;
      case ST_FIELD_BLOCK:
       next_state=parse_field(cur_class,tokens,lineno);
       break;
     }

     if (state!=next_state) {
      if (next_state==ST_INIT) {
       classes.add(cur_class);
       cur_class=null;
      }
      state=next_state;
     }
    }

    line=b.readLine();
    lineno++;
   }
  } catch(Throwable t) {
   t.printStackTrace();
  }

  return classes;
 }

 public static void write_utf8(DataOutputStream dos,String s) throws Throwable {
   dos.writeByte(CONSTANT_Utf8);
   if (s==null) s="null";
   dos.writeShort(s.length());
   dos.write(s.getBytes());
 }

 public static void write_class(DataOutputStream dos,int idx) throws Throwable {
   dos.writeByte(CONSTANT_Class);
   dos.writeShort(idx);
 }
 
 public static void write_integer(DataOutputStream dos,int val) throws Throwable {
   dos.writeByte(CONSTANT_Integer);
   dos.writeInt(val);
 }

 public static byte[] gen_class(ClassInfo ci) {
  byte[] data=null;
   
  try {
   ByteArrayOutputStream baos=new ByteArrayOutputStream();
   DataOutputStream dos=new DataOutputStream(baos);

   dos.writeInt(CLASS_MAGIC);
   dos.writeShort(CLASS_MINOR);
   dos.writeShort(CLASS_MAJOR);

   //initial idx
   int cp_cls_idx=2;   
   //classname, supername (CONSTANT_Utf8, CONSTANT_Class)
   int cp_int_idx=cp_cls_idx+2*2;
   //interface name (CONSTANT_Utf8, CONSTANT_Class)
   int cp_met_idx=cp_int_idx+2*ci.interfaces.size();
   //method name and descriptor (CONSTANT_Utf8, CONSTANT_Utf8)
   int cp_fld_idx=cp_met_idx+2*ci.methods.size();
   //field name and type (CONSTANT_Utf8, CONSTANT_Utf8)
   int cp_total=cp_fld_idx;

   for(int i=0;i<ci.fields.size();i++) {
    FieldInfo fi=(FieldInfo)ci.fields.elementAt(i);

    if ((fi.access&(ACC_STATIC|ACC_FINAL))==(ACC_STATIC|ACC_FINAL)) {
     if (fi.type.equals("I")) cp_total++;
    }

    cp_total+=2;
   }

   dos.writeShort(cp_total);

   //CP entry at idx 1
   write_utf8(dos,"ConstantValue");

   int cur_idx=cp_cls_idx;
   //CP entry for class
   write_utf8(dos,ci.name);
   write_class(dos,cur_idx);
   cur_idx+=2;
   //CP entry for superclass
   write_utf8(dos,ci.supername);
   write_class(dos,cur_idx);
   cur_idx+=2;

   //CP entries for interfaces
   for(int i=0;i<ci.interfaces.size();i++) {
    String iname=(String)ci.interfaces.elementAt(i);
    write_utf8(dos,iname);
    write_class(dos,cur_idx);
    cur_idx+=2;   
   }   

   //CP entries for methods
   for(int i=0;i<ci.methods.size();i++) {
    MethodInfo mi=(MethodInfo)ci.methods.elementAt(i);
    write_utf8(dos,mi.name);
    write_utf8(dos,mi.desc);
    cur_idx+=2;
   }  

   //CP entries for fields
   for(int i=0;i<ci.fields.size();i++) {
    FieldInfo fi=(FieldInfo)ci.fields.elementAt(i);
    write_utf8(dos,fi.name);
    write_utf8(dos,fi.type);

    if ((fi.access&(ACC_STATIC|ACC_FINAL))==(ACC_STATIC|ACC_FINAL)) {
     if (fi.type.equals("I")) {
      write_integer(dos,fi.val);
      cur_idx++;
     }
    }

    cur_idx+=2;
   }  

   dos.writeShort(ci.access);
   //this_class
   dos.writeShort(cp_cls_idx+1);
   //super_class
   if (ci.supername!=null) {
    dos.writeShort(cp_cls_idx+3);
   } else {
    dos.writeShort(0);
   }

   //interfaces
   int cnt=ci.interfaces.size();
   dos.writeShort(cnt);
   for(int i=0;i<cnt;i++) {
    dos.writeShort(cp_int_idx+2*i+1);
   }

   //fields
   cnt=ci.fields.size();
   dos.writeShort(cnt);

   int idx=0;
   for(int i=0;i<cnt;i++) {
    FieldInfo fi=(FieldInfo)ci.fields.elementAt(i);
    dos.writeShort(fi.access);
    dos.writeShort(cp_fld_idx+idx);
    dos.writeShort(cp_fld_idx+idx+1);

    if ((fi.access&(ACC_STATIC|ACC_FINAL))==(ACC_STATIC|ACC_FINAL)) {
     if (fi.type.equals("I")) {
      //attr count
      dos.writeShort(1);

      //attr name idx
      dos.writeShort(1);
      //attr len
      dos.writeInt(2);
      //attr val idx
      dos.writeShort(cp_fld_idx+idx+2);
      idx++;
     } else {
      dos.writeShort(0);
     }
    } else {
     dos.writeShort(0);
    }

    idx+=2;
   }

   //methods
   cnt=ci.methods.size();
   dos.writeShort(cnt);
   for(int i=0;i<cnt;i++) {
    MethodInfo mi=(MethodInfo)ci.methods.elementAt(i);
    dos.writeShort(mi.access);
    dos.writeShort(cp_met_idx+2*i);
    dos.writeShort(cp_met_idx+2*i+1);
    dos.writeShort(0);
   }

   //attributes
   dos.writeShort(0);

   data=baos.toByteArray();
  } catch(Throwable t) {
   t.printStackTrace();
  }

  return data;
 }

 public static void write_file(String name,byte data[]) {
  name=STUBS_DIR+"\\"+name+".class";

  int idx=name.lastIndexOf("/");

  if (idx>0) {
   String dir=name.substring(0,idx);
   try {
    new File(dir).mkdirs();
   } catch(Throwable t) {
    error("creating directory "+dir);
   }
  }

  try {
   FileOutputStream fos=new FileOutputStream(name,false);
   fos.write(data);
   fos.close();

   System.out.println("  "+name);
  } catch(Throwable t) {
   error("saving file "+name);
  }
 }

 public static void main(String args[]) {
  try {
   if (args.length!=1) usage();

   String classdesc=args[0];

   System.out.println("- loading API description file: "+classdesc);
   Vector classes=load_classdesc(classdesc);

   System.out.println("- generating stubs");

   int cnt=classes.size();

   for(int i=0;i<cnt;i++) {
    ClassInfo ci=(ClassInfo)classes.elementAt(i);
    byte data[]=gen_class(ci);

    write_file(ci.name,data);
   }

   System.out.println(" generated "+cnt+" class stubs");
  } catch (Throwable e) {
   e.printStackTrace();
  }
 }
}