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) {