Add first test to implement Wireshark compiler

joushx · Nov 25, 2016 · 4ae163e · 4ae163e
1 parent 8403e7f
commit 4ae163e
Showing 1 changed file with 70 additions and 31 deletions.
diff --git a/shared/src/main/scala/io/kaitai/struct/WiresharkClassCompiler.scala b/shared/src/main/scala/io/kaitai/struct/WiresharkClassCompiler.scala
@@ -24,35 +24,35 @@ class WiresharkClassCompiler(topClass: ClassSpec, out: LanguageOutputWriter) ext
   var currentTable: String = ""
 
   override def compile: Unit = {
-    out.puts("digraph {")
-    out.inc
-    out.puts("rankdir=LR;")
-    out.puts("node [shape=plaintext];")
-
     compileClass(topClass)
-
-    links.foreach { case (t1, t2, style) =>
-      out.puts(s"$t1 -> $t2 [$style];")
-    }
-
-    out.dec
-    out.puts("}")
   }
 
   def compileClass(curClass: ClassSpec): Unit = {
     provider.nowClass = curClass
     val className = curClass.name
 
-    out.puts(s"subgraph cluster__${type2class(className)} {")
-    out.inc
-    out.puts("label=\"" + type2display(className) + "\";")
-    out.puts("graph[style=dotted];")
+    out.puts(s"""${type2class(className)}_proto = Proto("${type2class(className)}","${type2display(className)}")""")
     out.puts
+    out.puts(s"local f = ${type2class(className)}_proto.fields")
 
     // Sequence
     compileSeq(className, curClass)
+    out.puts
 
-    curClass.instances.foreach { case (instName, instSpec) =>
+    // Dissect function
+    out.puts(s"function ${type2class(className)}_proto.dissector(buffer,pinfo,tree)")
+    out.inc
+    out.puts(s"""pinfo.cols.protocol = "${type2display(className)}"""")
+    out.puts(s"""main = tree:add(${type2display(className)}_proto, "${type2display(className)}")""")
+    out.puts
+
+    compileSeqFunction(className, curClass)
+
+    out.dec
+    out.puts("end")
+    out.puts
+
+    /*curClass.instances.foreach { case (instName, instSpec) =>
       instSpec match {
         case pis: ParseInstanceSpec =>
           tableStart(className, s"inst__${instName.name}")
@@ -61,36 +61,75 @@ class WiresharkClassCompiler(topClass: ClassSpec, out: LanguageOutputWriter) ext
         case vis: ValueInstanceSpec =>
           tableValueInstance(className, instName.name, vis)
       }
-    }
+    }*/
 
     //    curClass.enums.foreach { case(enumName, enumColl) => compileEnum(enumName, enumColl) }
 
-    out.add(extraClusterLines)
+    /*out.add(extraClusterLines)
     extraClusterLines.clear()
 
     // Recursive types
-    curClass.types.foreach { case (typeName, intClass) => compileClass(intClass) }
+    curClass.types.foreach { case (typeName, intClass) => compileClass(intClass) }*/
 
-    out.dec
-    out.puts("}")
+    out.puts(s"""${type2class(className)}_table = DissectorTable.get("tcp.port")""")
+    out.puts(s"tcp_table:add(<port>, ${type2class(className)}_proto)")
   }
 
   def compileSeq(className: List[String], curClass: ClassSpec): Unit = {
-    tableStart(className, "seq")
-    var seqPos: Option[Int] = Some(0)
     curClass.seq.foreach { (attr) =>
+
+      out.puts(s"\n-- ${attr.dataType.getClass.getName}")
+
       attr.id match {
         case NamedIdentifier(name) =>
-          tableRow(className, seqPos.map(_.toString), attr, name)
+          attr.dataType match{
+            case t: Int1Type => out.puts(s"""f.$name = ProtoField.uint8("${type2class(className)}.$name", "$name")""")
+            case t: IntMultiType =>
+
+              if(t.width.width == 2){
+                out.puts(s"""f.$name = ProtoField.uint16("${type2class(className)}.$name", "$name")""")
+              }
+            case t: BytesLimitType => out.puts(s"""f.$name = ProtoField.bytes("${type2class(className)}.$name", "$name")""")
+            case t: EnumType =>
+              out.puts(s"""f.$name = ProtoField.uint8("${type2class(className)}.$name", "$name")""") // TODO
+            case t: UserType =>
+              // do nothing
+            case _ => // TODO
+          }
+
+      }
+    }
+  }
 
-          val size = dataTypeSize(attr.dataType)
-          seqPos = (seqPos, size) match {
-            case (Some(pos), Some(siz)) => Some(pos + siz)
-            case _ => None
+  def compileSeqFunction(className: List[String], curClass: ClassSpec): Unit = {
+
+    out.puts(s"local offset = 0")
+
+    curClass.seq.foreach { (attr) =>
+
+      attr.id match {
+        case NamedIdentifier(name) =>
+          attr.dataType match{
+            case t: Int1Type =>
+              out.puts(s"main:add(buffer(offset,1), f.$name)")
+              out.puts("offset += 1")
+            case t: IntMultiType =>
+                out.puts(s"main:add(buffer(offset,${t.width.width}), f.$name)")
+                out.puts(s"offset += ${t.width.width}")
+            case t: BytesLimitType =>
+              out.puts(s"main:add(buffer(offset,X), f.$name)")
+              out.puts(s"offset += X")
+            case t: EnumType =>
+              out.puts(s"-- TODO") // TODO
+            case t: UserType =>
+              out.puts(s"""local ${t.name.head} = main:add("${t.name.head}")""")
+
+              // TODO print fields of seq
+            case _ => // TODO
           }
+
       }
     }
-    tableEnd
   }
 
   def compileParseInstance(className: List[String], id: InstanceIdentifier, inst: ParseInstanceSpec): Unit = {
@@ -360,7 +399,7 @@ class WiresharkClassCompiler(topClass: ClassSpec, out: LanguageOutputWriter) ext
 
 object WiresharkClassCompiler extends LanguageCompilerStatic {
   override def indent: String = "\t"
-  override def outFileName(topClassName: String): String = s"$topClassName.dot"
+  override def outFileName(topClassName: String): String = s"$topClassName.lua"
   override def getTranslator(tp: TypeProvider): BaseTranslator = new RubyTranslator(tp)
 
   def type2class(name: List[String]) = name.last