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");