1+ package net .bytebutcher .burpsendtoextension .builder ;
2+
3+ import com .google .common .collect .Lists ;
4+ import com .google .common .collect .Maps ;
5+ import net .bytebutcher .burpsendtoextension .models .CommandObject ;
6+ import net .bytebutcher .burpsendtoextension .models .Context ;
7+ import net .bytebutcher .burpsendtoextension .models .placeholder .IPlaceholderParser ;
8+ import net .bytebutcher .burpsendtoextension .models .placeholder .behaviour .CommandSeparatedPlaceholderBehaviour ;
9+ import net .bytebutcher .burpsendtoextension .models .placeholder .behaviour .FileSeparatedPlaceholderBehaviour ;
10+ import net .bytebutcher .burpsendtoextension .models .placeholder .behaviour .IPlaceholderBehaviour ;
11+ import net .bytebutcher .burpsendtoextension .models .placeholder .behaviour .StringSeparatedPlaceholderBehaviour ;
12+
13+ import java .io .File ;
14+ import java .io .PrintWriter ;
15+ import java .util .Comparator ;
16+ import java .util .List ;
17+ import java .util .Map ;
18+ import java .util .Objects ;
19+ import java .util .stream .Collectors ;
20+
21+ public class CommandBuilder {
22+
23+ // The model of the context menu entry with the format string and various options.
24+ private final CommandObject commandObject ;
25+
26+ // List of selected messages containing the placeholders and their values.
27+ private final List <Map <String , IPlaceholderParser >> placeholderMap ;
28+
29+ // List of placeholders and merged values.
30+ private final Map <Placeholder , String > placeholderValues ;
31+
32+ private final Context context ;
33+
34+ private static class Placeholder {
35+
36+ private final String name ;
37+ private final IPlaceholderBehaviour behaviour ;
38+
39+ public Placeholder (CommandObject .Placeholder placeholder ) {
40+ this .name = placeholder .getName ();
41+ this .behaviour = placeholder .getBehaviour ();
42+ }
43+
44+ @ Override
45+ public boolean equals (Object o ) {
46+ if (this == o ) return true ;
47+ if (o == null || getClass () != o .getClass ()) return false ;
48+ Placeholder that = (Placeholder ) o ;
49+ return Objects .equals (name , that .name ) && Objects .equals (behaviour , that .behaviour );
50+ }
51+
52+ @ Override
53+ public int hashCode () {
54+ return Objects .hash (name , behaviour );
55+ }
56+ }
57+
58+ public CommandBuilder (CommandObject commandObject , List <Map <String , IPlaceholderParser >> placeholderMap , Context context ) throws Exception {
59+ this .commandObject = commandObject ;
60+ this .placeholderMap = placeholderMap ;
61+ this .context = context ;
62+ this .placeholderValues = initPlaceholderValues ();
63+ }
64+
65+ public Map <Placeholder , String > initPlaceholderValues () throws Exception {
66+ Map <Placeholder , String > placeholderValues = Maps .newHashMap ();
67+ for (CommandObject .Placeholder coPlaceholder : commandObject .getPlaceholders ()) {
68+ Placeholder cbPlaceholder = new Placeholder (coPlaceholder );
69+ if (!placeholderValues .containsKey (cbPlaceholder )) {
70+ if (coPlaceholder .getBehaviour () instanceof StringSeparatedPlaceholderBehaviour ){
71+ // combine the values of all messages using the defined placeholder separator
72+ placeholderValues .put (cbPlaceholder , commandObject .getValid (placeholderMap , context ).stream ()
73+ .map (m -> m .get (coPlaceholder .getName ()))
74+ .map (iPlaceholder -> iPlaceholder .getValue (context ))
75+ .collect (Collectors .joining (((StringSeparatedPlaceholderBehaviour ) coPlaceholder .getBehaviour ()).getSeparator ())));
76+ } else if (coPlaceholder .getBehaviour () instanceof FileSeparatedPlaceholderBehaviour ){
77+ // combine the values of all messages and write them into a file.
78+ placeholderValues .put (cbPlaceholder , writeToFile (commandObject .getValid (placeholderMap , context ).stream ()
79+ .map (m -> m .get (coPlaceholder .getName ()))
80+ .map (iPlaceholder -> iPlaceholder .getValue (context ))
81+ .collect (Collectors .joining ("\n " ))));
82+ }
83+ }
84+ }
85+ return placeholderValues ;
86+ }
87+
88+ /**
89+ * Returns the command while all placeholders are replaced with their associated value as String.
90+ * @throws Exception when retrieving/replacing a placeholder failed.
91+ */
92+ public String build () throws Exception {
93+ try {
94+ List <String > result = Lists .newArrayList ();
95+ boolean containsCommandSeparatedPlaceholderBehaviour = commandObject .getPlaceholders ().stream ()
96+ .map (CommandObject .Placeholder ::getBehaviour )
97+ .anyMatch (c -> c instanceof CommandSeparatedPlaceholderBehaviour );
98+ if (containsCommandSeparatedPlaceholderBehaviour ) {
99+ for (int messageIndex = 0 ; messageIndex < placeholderMap .size (); messageIndex ++) {
100+ result .add (buildByMessage (messageIndex ));
101+ }
102+ } else {
103+ result .add (buildByMessage (0 ));
104+ }
105+ return String .join ("\n " , result );
106+ } catch (RuntimeException e ) {
107+ // Rethrow from unchecked to checked exception. We only deal with RuntimeException here, since streams
108+ // (here: placeholderMap.stream()) does not handle checked exceptions well.
109+ throw new Exception (e );
110+ }
111+ }
112+
113+ private String buildByMessage (int messageIndex ) throws Exception {
114+ StringBuffer format = new StringBuffer (commandObject .getFormat ());
115+ // For each placeholder, starting from the placeholder at the very end, replace it with the value from the message.
116+ List <CommandObject .Placeholder > placeholders = commandObject .getPlaceholders ().stream ().sorted (
117+ Comparator .comparing (CommandObject .Placeholder ::getEnd ).reversed ()
118+ ).collect (Collectors .toList ());
119+ for (CommandObject .Placeholder placeholder : placeholders ) {
120+ replaceCommandPlaceholder (placeholder , placeholderMap , messageIndex , format , context );
121+ }
122+ return format .toString ();
123+ }
124+
125+ private void replaceCommandPlaceholder (CommandObject .Placeholder placeholder , List <Map <String , IPlaceholderParser >> placeholderMap , int messageIndex , StringBuffer command , Context context ) throws Exception {
126+ String value ;
127+ if (placeholderValues .containsKey (new Placeholder (placeholder ))) {
128+ // merged values
129+ value = placeholderValues .get (new Placeholder (placeholder ));
130+ } else {
131+ // command separated - use the value from the actual message
132+ value = placeholderMap .get (messageIndex ).get (placeholder .getName ()).getValue (context );
133+ }
134+ command .replace (placeholder .getStart (), placeholder .getEnd (), value );
135+ }
136+
137+ private String writeToFile (String value ) throws Exception {
138+ try {
139+ File tmp = File .createTempFile ("burp_" , ".snd" );
140+ PrintWriter out = new PrintWriter (tmp .getPath ());
141+ out .write (value );
142+ out .flush ();
143+ return tmp .getAbsolutePath ();
144+ } catch (RuntimeException e ) {
145+ throw new Exception (this .getClass ().getSimpleName () + ": Error writing to temporary file!" , e );
146+ }
147+ }
148+ }
0 commit comments