snaprestore.m

   1 #import <Foundation/Foundation.h>
   2 #import <Foundation/NSFileManager.h>
   3 #import <IOKit/IOKitLib.h>
   4 #import <sys/snapshot.h>
   5 #import <getopt.h>
   6 #import "NSTask.h"
   7
   8 void usage(char *name) {
   9         printf(
  10                 "Usage: %s [volume] [snapshot]\n", name);
  11 }
  12
  13 NSString *bootsnapshot() {
  14         const io_registry_entry_t chosen = IORegistryEntryFromPath(0, "IODeviceTree:/chosen");
  15         const NSData *data = (__bridge const NSData *)IORegistryEntryCreateCFProperty(chosen, (__bridge CFStringRef)@"boot-manifest-hash", kCFAllocatorDefault, 0);
  16         IOObjectRelease(chosen);
  17
  18         NSMutableString *manifestHash = [NSMutableString stringWithString:@""];
  19         NSUInteger len = [data length];
  20         Byte *buf = (Byte*)malloc(len);
  21         memcpy(buf, [data bytes], len);
  22         int buf2;
  23         for (buf2 = 0; buf2 <= 19; buf2++) {
  24                 [manifestHash appendFormat:@"%02X", buf[buf2]];
  25         }
  26         // add com.apple.os.update-
  27         return [NSString stringWithFormat:@"%@%@", @"com.apple.os.update-", manifestHash];
  28 }
  29
  30 int restore(const char *vol, const char *snap) {
  31         int fd = open(vol, O_RDONLY, 0);
  32
  33         int ret = fs_snapshot_revert(fd, snap, 0);
  34         return ret;
  35 }
  36
  37 int mount(const char *vol, const char *snap, const char *mnt) {
  38         int fd = open(vol, O_RDONLY, 0);
  39
  40         BOOL isDir;
  41         NSFileManager *fileManager = [NSFileManager defaultManager];
  42                 if(![fileManager fileExistsAtPath:[NSString stringWithUTF8String:mnt] isDirectory:&isDir])
  43                         if(![fileManager createDirectoryAtPath:[NSString stringWithUTF8String:mnt] withIntermediateDirectories:YES attributes:nil error:NULL])
  44                                 NSLog(@"Error: Create folder failed %s", mnt);
  45
  46         int ret = fs_snapshot_mount(fd, mnt, snap, 0);
  47
  48         return ret;
  49 }
  50
  51 NSMutableSet *findApps(const char *root, const char *mnt) {
  52         NSMutableString *rootApplications = [NSMutableString stringWithUTF8String:root];
  53         rootApplications = [[rootApplications stringByAppendingString:@"/Applications"] mutableCopy];
  54
  55         NSMutableString *mntApplications = [NSMutableString stringWithUTF8String:mnt];
  56         mntApplications = [[mntApplications stringByAppendingString:@"/Applications"] mutableCopy];
  57
  58         NSArray *rootApps = [[NSFileManager defaultManager] contentsOfDirectoryAtPath:rootApplications error:nil];
  59         NSArray *mntApps = [[NSFileManager defaultManager] contentsOfDirectoryAtPath:mntApplications error:nil];
  60
  61         NSMutableSet *ret = [[NSMutableSet alloc] init];
  62         for (NSString *app in rootApps) {
  63                 if (![mntApps containsObject:app]) {
  64                         [ret addObject:[@"/Applications/" stringByAppendingString:app]];
  65                 }
  66         }
  67
  68         return ret;
  69 }
  70
  71 int rename(const char *vol, const char *snap) {
  72         int fd = open(vol, O_RDONLY, 0);
  73
  74         int ret = fs_snapshot_rename(fd, snap, [bootsnapshot() UTF8String], 0);
  75         return ret;
  76 }
  77
  78 int main(int argc, char *argv[]) {
  79         if (argc != 3) {
  80                 usage(argv[0]);
  81                 return 0;
  82         }
  83
  84         char *vol = argv[1];
  85         char *snap = argv[2];
  86         char *mnt = "/tmp/rootfsmnt";
  87
  88         printf("Restoring snapshot %s...\n", snap);
  89         restore(vol, snap);
  90         printf("Restored snapshot...\n");
  91         printf("Mounting rootfs...\n");
  92         mount(vol, snap, mnt);
  93         printf("Mounted %s at %s\n", snap, mnt);
  94         NSMutableSet *appSet = findApps(vol, mnt);
  95         if ([appSet count]) {
  96                 printf("Refreshing icon cache...\n");
  97                 NSMutableArray *argArray = [[NSMutableArray alloc] init];
  98                 for (NSString *app in appSet) {
  99                         [argArray addObject:@"-u"];
 100                         [argArray addObject:app];
 101                 }
 102                 NSTask *task = [[NSTask alloc] init];
 103                 [task setLaunchPath:@"/usr/bin/uicache"];
 104                 [task setArguments:argArray];
 105                 [task launch];
 106                 [task waitUntilExit];
 107         }
 108         printf("Renaming snapshot...\n");
 109         rename(vol, snap);
 110         printf("Restoring %s on %s has succeeded\n", snap, vol);
 111         return 0;
 112 }