1 files changed, 112 insertions, 0 deletions
diff --git a/snaprestore.m b/snaprestore.m
new file mode 100644
index 0000000..96e7b4d
--- /dev/null
+++ b/snaprestore.m
@@ -0,0 +1,112 @@
+#import <Foundation/Foundation.h>
+#import <Foundation/NSFileManager.h>
+#import <IOKit/IOKitLib.h>
+#import <sys/snapshot.h>
+#import <getopt.h>
+#import "NSTask.h"
+
+void usage(char *name) {
+	printf(
+		"Usage: %s [volume] [snapshot]\n", name);
+}
+
+NSString *bootsnapshot() {
+	const io_registry_entry_t chosen = IORegistryEntryFromPath(0, "IODeviceTree:/chosen");
+	const NSData *data = (__bridge const NSData *)IORegistryEntryCreateCFProperty(chosen, (__bridge CFStringRef)@"boot-manifest-hash", kCFAllocatorDefault, 0);
+	IOObjectRelease(chosen);
+
+	NSMutableString *manifestHash = [NSMutableString stringWithString:@""]; 
+	NSUInteger len = [data length];
+	Byte *buf = (Byte*)malloc(len);
+	memcpy(buf, [data bytes], len);
+	int buf2;
+	for (buf2 = 0; buf2 <= 19; buf2++) {
+		[manifestHash appendFormat:@"%02X", buf[buf2]];
+	}
+	// add com.apple.os.update-
+	return [NSString stringWithFormat:@"%@%@", @"com.apple.os.update-", manifestHash];
+}
+
+int restore(const char *vol, const char *snap) {
+	int fd = open(vol, O_RDONLY, 0);
+
+	int ret = fs_snapshot_revert(fd, snap, 0);
+	return ret;
+}
+
+int mount(const char *vol, const char *snap, const char *mnt) {
+	int fd = open(vol, O_RDONLY, 0);
+
+	BOOL isDir;
+	NSFileManager *fileManager = [NSFileManager defaultManager]; 
+		if(![fileManager fileExistsAtPath:[NSString stringWithUTF8String:mnt] isDirectory:&isDir])
+			if(![fileManager createDirectoryAtPath:[NSString stringWithUTF8String:mnt] withIntermediateDirectories:YES attributes:nil error:NULL])
+				NSLog(@"Error: Create folder failed %s", mnt);
+
+	int ret = fs_snapshot_mount(fd, mnt, snap, 0);
+	
+	return ret;
+}
+
+NSMutableSet *findApps(const char *root, const char *mnt) {
+	NSMutableString *rootApplications = [NSMutableString stringWithUTF8String:root];
+	rootApplications = [[rootApplications stringByAppendingString:@"/Applications"] mutableCopy];
+
+	NSMutableString *mntApplications = [NSMutableString stringWithUTF8String:mnt];
+	mntApplications = [[mntApplications stringByAppendingString:@"/Applications"] mutableCopy];
+
+	NSArray *rootApps = [[NSFileManager defaultManager] contentsOfDirectoryAtPath:rootApplications error:nil];
+	NSArray *mntApps = [[NSFileManager defaultManager] contentsOfDirectoryAtPath:mntApplications error:nil];
+
+	NSMutableSet *ret = [[NSMutableSet alloc] init];
+	for (NSString *app in rootApps) {
+		if (![mntApps containsObject:app]) {
+			[ret addObject:[@"/Applications/" stringByAppendingString:app]];
+		}
+	}
+
+	return ret;
+}
+
+int rename(const char *vol, const char *snap) {
+	int fd = open(vol, O_RDONLY, 0);
+
+	int ret = fs_snapshot_rename(fd, snap, [bootsnapshot() UTF8String], 0);
+	return ret;
+}
+
+int main(int argc, char *argv[]) {
+	if (argc != 3) {
+		usage(argv[0]);
+		return 0;
+	}
+
+	char *vol = argv[1];
+	char *snap = argv[2];
+	char *mnt = "/tmp/rootfsmnt";
+
+	printf("Restoring snapshot %s...\n", snap);
+	restore(vol, snap);
+	printf("Restored snapshot...\n");
+	printf("Mounting rootfs...\n");
+	mount(vol, snap, mnt);
+	printf("Mounted %s at %s\n", snap, mnt);
+	NSMutableSet *appSet = findApps(vol, mnt);
+	if (appSet) {
+		printf("Refreshing icon cache...\n");
+		NSMutableArray *argArray = [[NSMutableArray alloc] init];
+		for (NSString *app in appSet) {
+			[argArray addObject:@"-u"];
+			[argArray addObject:app];
+		}
+		NSTask *task = [[NSTask alloc] init];
+		[task setLaunchPath:@"/usr/bin/uicache"];
+		[task setArguments:argArray];
+		[task launch];
+		[task waitUntilExit];
+	}
+	printf("Renaming snapshot...\n");
+	rename(vol, snap);
+	printf("Restoring %s on %s has succeeded\n", snap, vol); 
+	return 0;
+}