idevicescreenshot: Choose a better filename, prevent overwriting existing files

1 files changed, 57 insertions, 18 deletions

diff --git a/tools/idevicescreenshot.c b/tools/idevicescreenshot.c
index 75a69ce..f33e83e 100644
--- a/tools/idevicescreenshot.c
+++ b/tools/idevicescreenshot.c
@@ -29,7 +29,9 @@
 #include <string.h>
 #include <stdlib.h>
 #include <errno.h>
+#include <math.h>
 #include <time.h>
+#include <unistd.h>
 #ifndef WIN32
 #include <signal.h>
 #endif
@@ -38,6 +40,7 @@
 #include <libimobiledevice/lockdown.h>
 #include <libimobiledevice/screenshotr.h>
 
+void get_image_filename(char *imgdata, char **filename);
 void print_usage(int argc, char **argv);
 
 int main(int argc, char **argv)
@@ -117,21 +120,7 @@ int main(int argc, char **argv)
 			char *imgdata = NULL;
 			uint64_t imgsize = 0;
 			if (screenshotr_take_screenshot(shotr, &imgdata, &imgsize) == SCREENSHOTR_E_SUCCESS) {
-				if (!filename) {
-					const char *fileext = NULL;
-					if (memcmp(imgdata, "\x89PNG", 4) == 0) {
-						fileext = ".png";
-					} else if (memcmp(imgdata, "MM\x00*", 4) == 0) {
-						fileext = ".tiff";
-					} else {
-						printf("WARNING: screenshot data has unexpected image format.\n");
-						fileext = ".dat";
-					}
-					time_t now = time(NULL);
-					filename = (char*)malloc(36);
-					size_t pos = strftime(filename, 36, "screenshot-%Y-%m-%d-%H-%M-%S", gmtime(&now));
-					sprintf(filename+pos, "%s", fileext);
-				}
+				get_image_filename(imgdata, &filename);
 				FILE *f = fopen(filename, "wb");
 				if (f) {
 					if (fwrite(imgdata, 1, (size_t)imgsize, f) == (size_t)imgsize) {
@@ -162,6 +151,53 @@ int main(int argc, char **argv)
 	return result;
 }
 
+void get_image_filename(char *imgdata, char **filename)
+{
+	// If the provided filename already has an extension, use it as is.
+	if (*filename) {
+		char *last_dot = strrchr(*filename, '.');
+		if (last_dot && !strchr(last_dot, '/')) {
+			return;
+		}
+	}
+
+	// Find the appropriate file extension for the filename.
+	const char *fileext = NULL;
+	if (memcmp(imgdata, "\x89PNG", 4) == 0) {
+		fileext = ".png";
+	} else if (memcmp(imgdata, "MM\x00*", 4) == 0) {
+		fileext = ".tiff";
+	} else {
+		printf("WARNING: screenshot data has unexpected image format.\n");
+		fileext = ".dat";
+	}
+
+	// If a filename without an extension is provided, append the extension.
+	// Otherwise, generate a filename based on the current time.
+	char *basename = NULL;
+	if (*filename) {
+		basename = (char*)malloc(strlen(*filename) + 1);
+		strcpy(basename, *filename);
+		free(*filename);
+	} else {
+		time_t now = time(NULL);
+		basename = (char*)malloc(32);
+		strftime(basename, 31, "screenshot-%Y-%m-%d-%H-%M-%S", gmtime(&now));
+	}
+
+	// Ensure the filename is unique on disk.
+	char *unique_filename = (char*)malloc(strlen(basename) + strlen(fileext) + 1);
+	sprintf(unique_filename, "%s%s", basename, fileext);
+	int i;
+	for (i = 2; access(unique_filename, F_OK) != -1; i++) {
+		free(unique_filename);
+		unique_filename = (char*)malloc(strlen(basename) + strlen(fileext) + floor(log10(i)) + 3);
+		sprintf(unique_filename, "%s-%d%s", basename, i, fileext);
+	}
+	*filename = unique_filename;
+	free(basename);
+}
+
 void print_usage(int argc, char **argv)
 {
 	char *name = NULL;
@@ -169,10 +205,13 @@ void print_usage(int argc, char **argv)
 	name = strrchr(argv[0], '/');
 	printf("Usage: %s [OPTIONS] [FILE]\n", (name ? name + 1: argv[0]));
 	printf("\n");
-	printf("Gets a screenshot from a device.\n");
+	printf("Gets a screenshot from a connected device.\n");
 	printf("\n");
-	printf("The screenshot is saved as a TIFF image with the given FILE name,\n");
-	printf("where the default name is \"screenshot-DATE.tiff\", e.g.:\n");
+	printf("The image is in PNG format for iOS 9+ and otherwise in TIFF format.\n");
+	printf("The screenshot is saved as an image with the given FILE name.\n");
+	printf("If FILE has no extension, FILE will be a prefix of the saved filename.\n");
+	printf("If FILE is not specified, \"screenshot-DATE\", will be used as a prefix\n");
+	printf("of the filename, e.g.:\n");
 	printf("   ./screenshot-2013-12-31-23-59-59.tiff\n");
 	printf("\n");
 	printf("NOTE: A mounted developer disk image is required on the device, otherwise\n");