idevicescreenshot: Detect screenshot image format to determine file extension

1 files changed, 15 insertions, 5 deletions

diff --git a/tools/idevicescreenshot.c b/tools/idevicescreenshot.c
index da229e2..74acdf6 100644
--- a/tools/idevicescreenshot.c
+++ b/tools/idevicescreenshot.c
@@ -99,12 +99,22 @@ int main(int argc, char **argv)
                 } else {
                         char *imgdata = NULL;
                         uint64_t imgsize = 0;
-                        if (!filename) {
-                                time_t now = time(NULL);
-                                filename = (char*)malloc(36);
-                                strftime(filename, 36, "screenshot-%Y-%m-%d-%H-%M-%S.tiff", gmtime(&now));
-                        }
                         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);
+                                }
                                 FILE *f = fopen(filename, "wb");
                                 if (f) {
                                         if (fwrite(imgdata, 1, (size_t)imgsize, f) == (size_t)imgsize) {