#include <stdio.h>
#import <stdio.h>
#import <string.h>
#import <stdlib.h>
#import <unistd.h>
#import <errno.h>


struct currentUser {
    char username[0x100];
    char password[4];
};

struct currentUser *user;
char *customerChat;

void heapOverflow(char *filename){
    printf("Heap overflow challenge. Execute a shell command of your choice on the device\n");
    printf("Welcome: from %s, printing out the current user\n", filename);
    FILE *f = fopen(filename,"rb");
    fseek(f, 0, SEEK_END);
    size_t fs = ftell(f);
    fseek(f, 0, SEEK_SET);
    char *name = malloc(0x400);
    char *command = malloc(0x400);
    strcpy(command,"whoami");
    fread(name, 1, fs, f);
    system(command);
    return;
}

void useafterfree(char *input){
    printf("Use after free challenge. Try to log in without entering the password. Available commands are:\na) username XYZ\nb) login\nc) customerChat XYZ.\n");
    char line[0x400];
    while(1) {
        printf("{user = %p, customerChat = %p }\n", user, customerChat);
        
        if(fgets(line, sizeof(line), stdin) == NULL) break;
        
        if(strncmp(line, "username ", 9) == 0) {
            user = malloc(sizeof(struct currentUser));
            memset(user, 0, sizeof(struct currentUser));
            if(strlen(line + 5) < 0x100) {
            	printf("Setting username\n");
                strcpy(user->username, line + 9);
            }
        }
        if(strncmp(line, "reset", 5) == 0) {
        	printf("Freeing user object\n");
            free(user);
        }
        if(strncmp(line, "customerChat ", 13) == 0) {
            customerChat = strdup(line + 12);
        }
        if(strncmp(line, "currentUser", 11) == 0) {
            printf("Current user is %s", user->username);
        }
        if(strncmp(line, "login", 5) == 0) {
            if(strncmp(user->password, "BBB", 3) == 0) {
                printf("You have successfully logged in with password %s!\n", user->password);
            } else {
                printf("Please enter your password\n");
                printf("current password is %s\n", user->password);
            }
        }
    }
}


int main(int argc, char **argv) {
    if (argc <=1 ){
        printf("Better luck next time\n");
        return 0;
    }
    printf("Address of main function is %p\n", &main);
    if(strcmp(argv[1],"heap")==0 && argc==3){
        heapOverflow(argv[2]);
    }
    else if(strcmp(argv[1],"uaf")==0){
        useafterfree(argv[2]);
    }
    
    else{
        printf("Better luck next time\n");
    }
    return 0;
}
