From 6c0766b4b4ab738fbcb588962a8ca3d3111bf84e Mon Sep 17 00:00:00 2001
From: wh1t3p1g <wh1t3p1g@gmail.com>
Date: Sun, 15 Oct 2023 12:39:42 +0800
Subject: [PATCH] add dump function

---
 cli/src/main/java/ysomap/cli/Console.java     | 33 ++++++++++++++-----
 cli/src/main/java/ysomap/cli/Session.java     |  5 +++
 .../java/ysomap/payloads/AbstractPayload.java | 11 +++++++
 .../main/java/ysomap/payloads/Payload.java    |  2 ++
 4 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/cli/src/main/java/ysomap/cli/Console.java b/cli/src/main/java/ysomap/cli/Console.java
index 2bef1ef..dbecfcd 100644
--- a/cli/src/main/java/ysomap/cli/Console.java
+++ b/cli/src/main/java/ysomap/cli/Console.java
@@ -146,7 +146,6 @@ public void dispatch(List<String> words) throws Exception {
                 break;
             case "dump":
                 dump();
-                Logger.success("Dump to dumped.yso.");
                 break;
             case "":
                 //do nothing
@@ -378,24 +377,40 @@ public void list() throws ArgumentsMissMatchException {
     }
 
     public void dump(){
+        if(args.isEmpty()){
+            args.add("dumped.yso");
+        }
+
         if(args.size() == 1 && curSession != null){
             StringBuilder sb = new StringBuilder();
-            Map<String, String> parameters = new HashMap<>();
             if(curSession.exploit != null){
+                sb.append("# exploit settings\n");
                 sb.append(curSession.exploit.dump());
-                parameters.putAll(curSession.exploit.getAllParameters());
+                Map<String, String> parameters = curSession.exploit.getAllParameters();
+
+                for(Map.Entry<String, String> entry:parameters.entrySet()){
+                    sb.append(String.format("set %s %s\n", entry.getKey(), entry.getValue()));
+                }
             }
 
             if(curSession.payload != null){
                 if(curSession.bullet != null){
                     curSession.payload.setBullet(curSession.bullet);
                 }
+                sb.append("\n");
+                sb.append("# payload settings\n");
                 sb.append(curSession.payload.dump());
-                parameters.putAll(curSession.payload.getAllParameters());
-            }
-
-            for(Map.Entry<String, String> entry:parameters.entrySet()){
-                sb.append(String.format("set %s %s\n", entry.getKey(), entry.getValue()));
+                sb.append("# serializer settings\n");
+                sb.append(String.format("set serializeType %s\n", curSession.payload.getSerializeType()));
+                sb.append(String.format("set encoder %s\n", curSession.payload.getEncoder()));
+                sb.append(String.format("set output %s\n", curSession.payload.getOutputType()));
+                sb.append(String.format("set serialVersionUID %s\n", curSession.payload.getSerialVersionUID()));
+                sb.append(String.format("set checkRunning %s\n", curSession.payload.getCheckRunning()));
+                Map<String, String> parameters = curSession.payload.getAllParameters();
+                sb.append("# bullet settings\n");
+                for(Map.Entry<String, String> entry:parameters.entrySet()){
+                    sb.append(String.format("set %s %s\n", entry.getKey(), entry.getValue()));
+                }
             }
 
             if(sb.toString().isEmpty()){
@@ -403,9 +418,11 @@ public void dump(){
                 return;
             }
 
+            sb.append("# start to run\n");
             sb.append("run\n");
             try {
                 FileHelper.filePutContent(args.get(0), sb.toString().getBytes());
+                Logger.success(String.format("Dump settings to '%s'.", args.get(0)));
             } catch (IOException e) {
                 throw new RuntimeException(e);
             }
diff --git a/cli/src/main/java/ysomap/cli/Session.java b/cli/src/main/java/ysomap/cli/Session.java
index 7ee55a4..aca43f2 100644
--- a/cli/src/main/java/ysomap/cli/Session.java
+++ b/cli/src/main/java/ysomap/cli/Session.java
@@ -123,6 +123,10 @@ public HashMap<String, Object> update(Class<?> clazz){
     }
 
     public void setValue(String key, String value) throws ArgumentsMissMatchException {
+        if("null".equals(value)){
+            value = null;
+        }
+
         if("serializeType".equals(key) && payload != null){
             payload.setSerializeType(value);
             return;
@@ -137,6 +141,7 @@ public void setValue(String key, String value) throws ArgumentsMissMatchExceptio
             return;
         }else if("checkRunning".equals(key)){
             isCheckRunning = Boolean.parseBoolean(value);
+            payload.setCheckRunning(isCheckRunning);
             return;
         }
 
diff --git a/core/src/main/java/ysomap/payloads/AbstractPayload.java b/core/src/main/java/ysomap/payloads/AbstractPayload.java
index 1466f10..8c5aae7 100644
--- a/core/src/main/java/ysomap/payloads/AbstractPayload.java
+++ b/core/src/main/java/ysomap/payloads/AbstractPayload.java
@@ -24,6 +24,7 @@ public abstract class AbstractPayload<T> implements Payload<T>{
     // 自由定义
     public String serializerEncoder;
     public String serializerSerialVersionUID;
+    public boolean checkRunning = false;
     public Bullet bullet;
 
     public AbstractPayload(){
@@ -54,6 +55,16 @@ public void setSerialVersionUID(String uid) {
         this.serializerSerialVersionUID = uid;
     }
 
+    @Override
+    public void setCheckRunning(boolean flag) {
+        checkRunning = flag;
+    }
+
+    @Override
+    public boolean getCheckRunning() {
+        return checkRunning;
+    }
+
     @Override
     final public T getObject() throws Exception {
         // check bullet args first
diff --git a/core/src/main/java/ysomap/payloads/Payload.java b/core/src/main/java/ysomap/payloads/Payload.java
index 48f9f91..f9939c1 100644
--- a/core/src/main/java/ysomap/payloads/Payload.java
+++ b/core/src/main/java/ysomap/payloads/Payload.java
@@ -28,7 +28,9 @@ public interface Payload<T> {
     void setOutputType(String outputType);
 
     void setSerialVersionUID(String uid);
+    void setCheckRunning(boolean flag);
 
+    boolean getCheckRunning();
     String getSerializeType();
 
     String getEncoder();