#include <unistd.h>
#include <sys/syscall.h>
#include <fcntl.h>
#include <string.h>
#include <stdlib.h>
#include <stdint.h>

// Include yescrypt implementation
#include "yescrypt.h"

#define BUF_SIZE 4096
#define MAX_LINE 1024
#define MAX_PASS 256

// Minimal strlen
static int my_strlen(const char *s) { int n=0; while(s[n]) n++; return n; }

// Minimal strncmp
static int my_strncmp(const char *a, const char *b, int n) { for(int i=0;i<n;i++) if(a[i]!=b[i]||!a[i]||!b[i]) return a[i]-b[i]; return 0; }

// Write string to stderr
static void write_stderr(const char *s) {
    syscall(SYS_write, 2, s, my_strlen(s));
}

// Write string to stdout
static void write_stdout(const char *s) {
    syscall(SYS_write, 1, s, my_strlen(s));
}

// Find position of substring in string
static int strpos(const char *string, const char *substring, int offset) {
    int string_len = my_strlen(string);
    int sub_len = my_strlen(substring);
    
    for(int i = offset; i <= string_len - sub_len; i++) {
        int j;
        for(j = 0; j < sub_len; j++) {
            if(string[i + j] != substring[j]) break;
        }
        if(j == sub_len) return i;
    }
    return -1;
}

// Read a line from fd into buf, return length or 0 on EOF
static int read_line(int fd, char *buf, int max) {
    int i=0, r;
    char c;
    while(i<max-1) {
        r = syscall(SYS_read, fd, &c, 1);
        if(r<=0) break;
        if(c=='\n') break;
        buf[i++] = c;
    }
    buf[i]=0;
    return i;
}

// Extract user hash from shadow line
static char* extract_user_hash(const char *shadow_line, const char *username) {
    int username_index = strpos(shadow_line, username, 0);
    if(username_index == -1) return NULL;
    
    // Find the start of hash (after username:)
    int hash_start = username_index + my_strlen(username) + 1;
    int hash_end = strpos(shadow_line, ":", hash_start);
    if(hash_end == -1) return NULL;
    
    int hash_len = hash_end - hash_start;
    char *user_hash = malloc(hash_len + 1);
    memcpy(user_hash, shadow_line + hash_start, hash_len);
    user_hash[hash_len] = 0;
    
    return user_hash;
}

// Parse hash to extract algorithm ID, salt, and hash value
static void parse_hash(const char *user_hash, int *hash_id, char *salt, char *hash_value) {
    // Find positions of $ characters
    int first_dollar = strpos(user_hash, "$", 0);
    if(first_dollar == -1) return;
    
    int second_dollar = strpos(user_hash, "$", first_dollar + 1);
    if(second_dollar == -1) return;
    
    int third_dollar = strpos(user_hash, "$", second_dollar + 1);
    if(third_dollar == -1) return;
    
    // Extract hash ID (algorithm)
    int id_len = second_dollar - first_dollar - 1;
    char id_str[8] = {0};
    memcpy(id_str, user_hash + first_dollar + 1, id_len);
    
    // Handle yescrypt (y) and numeric IDs
    if(id_len == 1 && id_str[0] == 'y') {
        *hash_id = 'y';  // yescrypt
        
        // For yescrypt, we need to find the 4th $ to get the actual salt
        int fourth_dollar = strpos(user_hash, "$", third_dollar + 1);
        if(fourth_dollar == -1) return;
        
        // Extract parameters (between 2nd and 3rd $)
        int param_len = third_dollar - second_dollar - 1;
        char parameters[64] = {0};
        memcpy(parameters, user_hash + second_dollar + 1, param_len);
        parameters[param_len] = 0;
        
        // Extract salt (between 3rd and 4th $)
        int salt_len = fourth_dollar - third_dollar - 1;
        memcpy(salt, user_hash + third_dollar + 1, salt_len);
        salt[salt_len] = 0;
        
        // Extract hash value (after 4th $)
        int hash_len = my_strlen(user_hash) - fourth_dollar - 1;
        memcpy(hash_value, user_hash + fourth_dollar + 1, hash_len);
        hash_value[hash_len] = 0;
        
        // For yescrypt, we need to include parameters in the salt for crypt()
        char combined_salt[128];
        int len = 0;
        combined_salt[len++] = '$';
        combined_salt[len++] = 'y';
        combined_salt[len++] = '$';
        memcpy(combined_salt + len, parameters, param_len);
        len += param_len;
        combined_salt[len++] = '$';
        memcpy(combined_salt + len, salt, salt_len);
        len += salt_len;
        combined_salt[len] = 0;
        
        // Copy the combined salt back to the salt parameter
        memcpy(salt, combined_salt, len + 1);
        
    } else {
        *hash_id = atoi(id_str);
        
        // Extract salt (standard format)
        int salt_len = third_dollar - second_dollar - 1;
        memcpy(salt, user_hash + second_dollar + 1, salt_len);
        salt[salt_len] = 0;
        
        // Extract hash value
        int hash_len = my_strlen(user_hash) - third_dollar - 1;
        memcpy(hash_value, user_hash + third_dollar + 1, hash_len);
        hash_value[hash_len] = 0;
    }
}

// Create formatted salt string for crypt()
static char* create_salt(int hash_id, const char *salt) {
    char *formatted_salt = malloc(128);
    int len = 0;
    
    if(hash_id == 'y') {
        // For yescrypt, the salt is already properly formatted
        int salt_len = my_strlen(salt);
        memcpy(formatted_salt, salt, salt_len);
        formatted_salt[salt_len] = 0;
    } else {
        // Standard format for other algorithms
        formatted_salt[len++] = '$';
        if(hash_id == 1) {
            formatted_salt[len++] = '1';
        } else if(hash_id == 5) {
            formatted_salt[len++] = '5';
        } else if(hash_id == 6) {
            formatted_salt[len++] = '6';
        } else {
            // Default to SHA512
            formatted_salt[len++] = '6';
        }
        formatted_salt[len++] = '$';
        
        int salt_len = my_strlen(salt);
        memcpy(formatted_salt + len, salt, salt_len);
        len += salt_len;
        formatted_salt[len] = 0;
    }
    
    return formatted_salt;
}

// Minimal itoa for positive integers
static void my_itoa(int value, char *buf) {
    char tmp[16];
    int i = 0, j = 0;
    if (value == 0) { buf[0] = '0'; buf[1] = 0; return; }
    while (value > 0) { tmp[i++] = '0' + (value % 10); value /= 10; }
    while (i > 0) { buf[j++] = tmp[--i]; }
    buf[j] = 0;
}

// Custom yescrypt crypt function using reference implementation parsing
static char* yescrypt_crypt(const char *password, const char *full_hash) {
    static char result[128];
    
    // Use the reference implementation's yescrypt function
    uint8_t *generated_hash = yescrypt((const uint8_t *)password, (const uint8_t *)full_hash);
    if (!generated_hash) {
        return NULL;
    }
    
    // Compare the generated hash with the expected hash
    if (strcmp((char *)generated_hash, full_hash) == 0) {
        // Copy the generated hash to our result buffer
        int len = strlen((char *)generated_hash);
        if (len < sizeof(result)) {
            memcpy(result, generated_hash, len + 1);
            return result;
        }
    }
    
    return NULL;
}

int main(int argc, char **argv) {
    if(argc<3) {
        write_stderr("Usage: ");
        write_stderr(argv[0]);
        write_stderr(" <username> <wordlist>\n");
        return 1;
    }
    char *username = argv[1];
    char *wordlist_path = argv[2];
    const char *shadow_path = "/etc/shadow";

    int sfd = syscall(SYS_open, shadow_path, O_RDONLY, 0);
    if(sfd<0) {
        write_stderr("Failed to open shadow file: ");
        write_stderr(shadow_path);
        write_stderr("\n");
        return 2;
    }
    
    char line[MAX_LINE];
    char *user_hash = NULL;
    while(read_line(sfd, line, MAX_LINE)) {
        user_hash = extract_user_hash(line, username);
        if(user_hash) break;
    }
    syscall(SYS_close, sfd);
    
    if(!user_hash) {
        write_stderr("User '");
        write_stderr(username);
        write_stderr("' not found in shadow file.\n");
        return 3;
    }

    // Parse the hash
    int hash_id;
    char salt[64] = {0};
    char hash_value[512] = {0};
    parse_hash(user_hash, &hash_id, salt, hash_value);
    
    if(hash_id == 0) {
        write_stderr("Failed to parse hash format for user '");
        write_stderr(username);
        write_stderr("'\n");
        free(user_hash);
        return 4;
    }
    
    // Print useful information
    write_stderr("Target user: ");
    write_stderr(username);
    write_stderr("\n");
    
    write_stderr("Hash type: ");
    if(hash_id == 'y') {
        write_stderr("yescrypt");
    } else if(hash_id == 1) {
        write_stderr("MD5");
    } else if(hash_id == 5) {
        write_stderr("SHA256");
    } else if(hash_id == 6) {
        write_stderr("SHA512");
    } else {
        char id_str[8];
        my_itoa(hash_id, id_str);
        write_stderr(id_str);
    }
    write_stderr("\n");
    
    write_stderr("Full hash: ");
    write_stderr(user_hash);
    write_stderr("\n");
    
    write_stderr("Starting bruteforce...\n");
    
    // Create formatted salt for crypt()
    char *formatted_salt = create_salt(hash_id, salt);
    
    int wfd = syscall(SYS_open, wordlist_path, O_RDONLY, 0);
    if(wfd<0) { 
        write_stderr("Failed to open wordlist: ");
        write_stderr(wordlist_path);
        write_stderr("\n");
        free(user_hash);
        free(formatted_salt);
        return 5; 
    }
    
    char pass[MAX_PASS];
    int found=0;
    char c;
    int n=0;
    int total_tried=0;
    
    while(syscall(SYS_read, wfd, &c, 1) == 1) {
        if(c == '\n' || c == '\r') {
            if(n > 0) {
                pass[n] = 0;
                total_tried++;
                
                // Show progress every 1000 attempts
                if(total_tried % 1000 == 0) {
                    char progress_str[32];
                    write_stderr("Tried ");
                    my_itoa(total_tried, progress_str);
                    write_stderr(progress_str);
                    write_stderr(" passwords...\n");
                }
                
                char *try_hash;
                
                if(hash_id == 'y') {
                    // Use yescrypt for yescrypt hashes
                    try_hash = yescrypt_crypt(pass, user_hash);
                } else {
                    // Use system crypt() for other hash types
                    try_hash = crypt(pass, formatted_salt);
                }
                
                if(try_hash && strcmp(try_hash, user_hash)==0) {
                    write_stdout("Found password: ");
                    write_stdout(pass);
                    write_stdout("\n");
                    found=1;
                    break;
                }
                n = 0;
            }
        } else if(n < MAX_PASS-1) {
            pass[n++] = c;
        }
    }
    
    // Handle last line if no newline at end
    if(n > 0 && !found) {
        pass[n] = 0;
        total_tried++;
        char *try_hash;
        
        if(hash_id == 'y') {
            // Use yescrypt for yescrypt hashes
            try_hash = yescrypt_crypt(pass, user_hash);
        } else {
            // Use system crypt() for other hash types
            try_hash = crypt(pass, formatted_salt);
        }
        
        if(try_hash && strcmp(try_hash, user_hash)==0) {
            write_stdout("Found password: ");
            write_stdout(pass);
            write_stdout("\n");
            found=1;
        }
    }
    
    syscall(SYS_close, wfd);
    
    // Print final summary
    char total_str[32];
    write_stderr("Total passwords tried: ");
    my_itoa(total_tried, total_str);
    write_stderr(total_str);
    write_stderr("\n");
    
    if(!found) {
        write_stderr("No password match found for user '");
        write_stderr(username);
        write_stderr("' in wordlist.\n");
    } else {
        write_stderr("Password successfully cracked!\n");
    }
    
    free(user_hash);
    free(formatted_salt);
    return found?0:6;
}