diff --git a/src/FileManager.java b/src/FileManager.java index c543b3b..aefa63f 100644 --- a/src/FileManager.java +++ b/src/FileManager.java @@ -1,8 +1,7 @@ -import java.io.File; -import java.io.FileWriter; -import java.io.IOException; +import java.io.*; import java.util.Scanner; +// TODO: close scanner etc after use public class FileManager { public static void writeFile(String fileLocation, String writeString) throws IOException { File newFile = new File(fileLocation); @@ -22,4 +21,16 @@ public class FileManager { return fileReader.useDelimiter("\\Z").next(); } } + + public static Object readObjectFromFile(String fileLocation, String allowedObject) throws IOException, ClassNotFoundException { + FileInputStream fileStream = new FileInputStream(fileLocation); + SafeObjectInputStream objectStream = new SafeObjectInputStream(fileStream, allowedObject); + return objectStream.readObject(); + } + + public static void writeObjectToFile(String fileLocation, Object o) throws IOException { + FileOutputStream fileStream = new FileOutputStream(fileLocation); + ObjectOutputStream objectStream = new ObjectOutputStream(fileStream); + objectStream.writeObject(o); + } } diff --git a/src/SafeObjectInputStream.java b/src/SafeObjectInputStream.java new file mode 100644 index 0000000..3045b32 --- /dev/null +++ b/src/SafeObjectInputStream.java @@ -0,0 +1,24 @@ +import java.io.IOException; +import java.io.InputStream; +import java.io.ObjectInputStream; +import java.io.ObjectStreamClass; + +public class SafeObjectInputStream extends ObjectInputStream { + + String allowedClass; + + public SafeObjectInputStream(InputStream in, String allowedClass) throws IOException { + super(in); + this.allowedClass = allowedClass; + } + + // local files are generally assumed to be safe, but this additional check slightly hardens the application against using arbitrary attacks to exploit + // please note that it is still insecure; don't use in sensitive contexts! + @Override + protected Class resolveClass(ObjectStreamClass desc) throws IOException, ClassNotFoundException { + if (!desc.getName().equals(allowedClass)) { + throw new SecurityException(); + } + return super.resolveClass(desc); + } +} diff --git a/src/UtilityFunction.java b/src/UtilityFunction.java index 7facb48..764d65e 100644 --- a/src/UtilityFunction.java +++ b/src/UtilityFunction.java @@ -1,4 +1,6 @@ import java.awt.event.KeyEvent; +import java.io.IOException; +import java.io.ObjectInputStream; import java.lang.reflect.Array; import java.util.ArrayList;