diff --git a/src/main.c b/src/main.c
index 728721f..93831af 100755
--- a/src/main.c
+++ b/src/main.c
@@ -102,6 +102,7 @@ int main(int argc, char *argv[]) {
     // Polybius square
     else if(strcmp(argv[1], "-s") == 0){
         char key[26];
+        int res;
 
         printf("Your key: ");
         scanf("%s", key);
@@ -112,9 +113,12 @@ int main(int argc, char *argv[]) {
         }
 
         if(strcmp(hook, "crypt") == 0)
-            cryptPolybius(fileSrc, fileDst, key);
+            res = cryptPolybius(fileSrc, fileDst, key);
         else
-            decryptPolybius(fileSrc, fileDst, key);
+            res = decryptPolybius(fileSrc, fileDst, key);
+
+        if (res != 0)
+            printf("Failed to crypt/decrypt the file\n");
     }
 
     /* Display error */
@@ -125,7 +129,5 @@ int main(int argc, char *argv[]) {
      */
 
     free(fileDst);
-    //free(buffer);
-    //free(data);
     return 0;
 }
diff --git a/src/polybius.c b/src/polybius.c
index 318c411..a1c9c60 100644
--- a/src/polybius.c
+++ b/src/polybius.c
@@ -42,8 +42,10 @@ int cryptPolybius(const char *filenameSrc, const char *filenameDst, const char *
                     }
                 }
             }
-            if (c == '\n')
-                fputc('\n', fDst);
+            if (c == '\n' || c == ' '){
+                fputc(c, fDst);
+                fputc(0, fDst); // For having a couple of number
+            }
         }
 
         size = 0; 
@@ -62,7 +64,6 @@ int decryptPolybius(const char *filenameSrc, const char *filenameDst, const char
     char data[BUFFER_SIZE];
     char dataEncrypted[BUFFER_SIZE];
     int size = 0, end = 0, i = 0, j = 0;
-    char c = 0;
 
     generateSquarePolybius(polybius, key);
 
@@ -85,17 +86,17 @@ int decryptPolybius(const char *filenameSrc, const char *filenameDst, const char
 
 	    /********* Decryption **********/
         for (i = 0; i < size; i += 2){
-            if (data[i] >= 0 && data[i] <= POLYBIUS_SIZE) {
+            if (data[i] >= 48 && data[i] <= 53) {
                 int c1 = data[i] - '0';
                 int c2 = data[i+1] - '0';
                 char c = polybius[c1][c2];
-                //printf("%d%d ", c1, c2);
-                //printf("%d ", c1);
                 printf("%c", c);
-                //printf("%c", data[i]);
+                fputc(c, fDst);
             }
-            else
+            else{
                 printf("%c", data[i]);
+                fputc(data[i], fDst);
+            }
         }
 
         size = 0;