package main

import (
    "encoding/json"
    "flag"
    "fmt"
    "net/http"
    "net/url"
    "os"
    "strings"
    "time"

    "github.com/go-git/go-git/v5"
    "github.com/go-git/go-git/v5/plumbing"
)

func main() {
    flag.Usage = func() {
        fmt.Fprintf(flag.CommandLine.Output(), "this program searches github for public repositories created today that match the keyword 'cve-<current_year>'\n")
        fmt.Fprintf(flag.CommandLine.Output(), "usage:\n")
        flag.PrintDefaults()
    }
    githubToken := flag.String("token", "", "github api token")
    cloneDir := flag.String("cloneDir", "cve-pocs", "directory to clone repositories")
    clonedListFile := flag.String("clonedList", "cve-pocs.txt", "file to store cloned repository urls")
    autoUpdate := flag.Bool("auto-update", false, "automatically update previously cloned repositories")
    customDate := flag.String("date", "", "specify a custom date in YYYY-MM-DD format")
    silent := flag.Bool("silent", false, "suppress update messages")
    flag.Parse()

    readmeFile := *cloneDir + "/README.md"

    if *githubToken == "" {
        fmt.Println("[err] github token is required; use -token flag to provide the token")
        return
    }

    if _, err := os.Stat(*cloneDir); os.IsNotExist(err) {
        err := os.Mkdir(*cloneDir, 0755)
        if err != nil {
            fmt.Printf("[err] failed to create clone directory: %v\n", err)
            return
        }
    }

    clonedRepos := loadClonedRepos(*clonedListFile)

    //today := time.Now().UTC().Format("2006-01-02")
    var today string
    if *customDate != "" {
        _, err := time.Parse("2006-01-02", *customDate)
        if err != nil {
            fmt.Printf("[err] invalid date format: %s; use YYYY-MM-DD\n", *customDate)
            return
        }
        today = *customDate
    } else {
        today = time.Now().UTC().Format("2006-01-02")
    }

    year := time.Now().Year()
    KEYWORD := fmt.Sprintf("cve-%d", year)

    fmt.Printf("[inf] searching for repositories with keyword: %s, created on: %s\n", KEYWORD, today)

    baseURL := "https://api.github.com"
    resource := "/search/repositories"

    params := url.Values{}
    params.Add("q", fmt.Sprintf("\"%s\" created:%s", KEYWORD, today))
    params.Add("sort", "updated")
    params.Add("order", "desc")
    params.Add("per_page", "100")

    u, err := url.ParseRequestURI(baseURL)
    if err != nil {
        fmt.Printf("[err] failed to parse url: %v\n", err)
        return
    }
    u.Path = resource
    u.RawQuery = params.Encode()

    urlStr := fmt.Sprintf("%v", u)
    fmt.Println("[inf] github api url:", urlStr)

    req, err := http.NewRequest("GET", urlStr, nil)
    if err != nil {
        fmt.Printf("[err] failed to create http request: %v\n", err)
        return
    }

    req.Header.Add("Authorization", "token "+*githubToken)
    req.Header.Add("Accept", "application/vnd.github.v3+json")

    client := &http.Client{}
    resp, err := client.Do(req)
    if err != nil {
        fmt.Printf("[err] failed to make GET request: %v\n", err)
        return
    }
    defer resp.Body.Close()

    if resp.StatusCode != http.StatusOK {
        fmt.Printf("[err] received non-ok http status %s\n", resp.Status)
        return
    }

    var result map[string]interface{}
    if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
        fmt.Printf("[err] failed to parse response: %v\n", err)
        return
    }

    repos, ok := result["items"].([]interface{})
    if !ok || len(repos) == 0 {
        fmt.Println("[wrn] no repositories found or failed to parse repositories")
        return
    }

    for _, repo := range repos {
        if repoMap, ok := repo.(map[string]interface{}); ok {
            repoName := repoMap["name"].(string)
            owner := repoMap["owner"].(map[string]interface{})["login"].(string)
            repoURL := repoMap["html_url"].(string)

            description, ok := repoMap["description"].(string)
            if !ok {
                description = "No description"
            }

            cloneName := fmt.Sprintf("%s_%s", owner, repoName)

            if _, cloned := clonedRepos[repoURL]; !cloned {
                fmt.Printf("[inf] new poc: %s\n", repoURL)
                fmt.Printf("[inf] description: %s\n", description)
                fmt.Println()

                if cloneRepo(repoURL, *cloneDir, cloneName) {
                    clonedRepos[repoURL] = struct{}{}
                    appendToFile(*clonedListFile, repoURL)
                    updateReadme(readmeFile, repoURL, description)
                }
            }
        }
    }
    if *autoUpdate {
        updateClonedRepositories(*cloneDir, clonedRepos, *silent)
    }
}

func cloneRepo(repoURL, cloneDir, cloneName string) bool {
    _, err := git.PlainClone(cloneDir+"/"+cloneName, false, &git.CloneOptions{
        URL: repoURL,
    })
    if err != nil {
        fmt.Printf("[err] failed to clone repository: %v\n", err)
        return false
    }
    return true
}

func loadClonedRepos(filePath string) map[string]struct{} {
    clonedRepos := make(map[string]struct{})

    file, err := os.Open(filePath)
    if err != nil {
        if os.IsNotExist(err) {
            return clonedRepos
        }
        fmt.Printf("[err] failed to open cloned repos list file: %v\n", err)
        return clonedRepos
    }
    defer file.Close()

    var line string
    for {
        _, err := fmt.Fscanln(file, &line)
        if err != nil {
            break
        }
        clonedRepos[strings.TrimSpace(line)] = struct{}{}
    }

    return clonedRepos
}

func appendToFile(filePath, line string) {
    file, err := os.OpenFile(filePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
    if err != nil {
        fmt.Printf("[err] failed to open file %s: %v\n", filePath, err)
        return
    }
    defer file.Close()

    if _, err := file.WriteString(line + "\n"); err != nil {
        fmt.Printf("[err] failed to write to file %s: %v\n", filePath, err)
    }
}

func updateReadme(readmeFile, repoURL, description string) {
    file, err := os.OpenFile(readmeFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
    if err != nil {
        fmt.Printf("[err] failed to open README file %s: %v\n", readmeFile, err)
        return
    }
    defer file.Close()

    entry := fmt.Sprintf("- [%s](%s) %s\n", repoURL, repoURL, description)
    if _, err := file.WriteString(entry); err != nil {
        fmt.Printf("[err] failed to write to README file %s: %v\n", readmeFile, err)
    }
}

func updateClonedRepositories(cloneDir string, clonedRepos map[string]struct{}, silent bool) {
    for repoURL := range clonedRepos {
        cloneName := extractRepoNameFromURL(repoURL)
        repoPath := fmt.Sprintf("%s/%s", cloneDir, cloneName)

        r, err := git.PlainOpen(repoPath)
        if err != nil {
            if !silent {
                fmt.Printf("[err] failed to open repository %s: %v\n", repoPath, err)
            }
            continue
        }

        headRef, err := r.Head()
        if err != nil {
            if !silent {
                fmt.Printf("[err] failed to get HEAD for repository %s: %v\n", repoPath, err)
            }
            continue
        }

        defaultBranch := ""
        if headRef.Name().IsBranch() {
            defaultBranch = headRef.Name().Short()
        } else {
            if !silent {
                fmt.Printf("[err] HEAD is not a branch for repository %s, skipping\n", repoPath)
            }
            continue
        }

        err = r.Fetch(&git.FetchOptions{
            RemoteName: "origin",
        })
        if err != nil && err != git.NoErrAlreadyUpToDate {
            if !silent {
                fmt.Printf("[err] failed to fetch updates for repository %s: %v\n", repoPath, err)
            }
            continue
        }

        w, err := r.Worktree()
        if err != nil {
            if !silent {
                fmt.Printf("[err] failed to get worktree for repository %s: %v\n", repoPath, err)
            }
            continue
        }

        remoteBranchRef := fmt.Sprintf("origin/%s", defaultBranch)
        remoteRef, err := r.Reference(plumbing.ReferenceName("refs/remotes/"+remoteBranchRef), true)
        if err != nil {
            if !silent {
                fmt.Printf("[err] failed to get reference for %s in repository %s: %v\n", remoteBranchRef, repoPath, err)
            }
            continue
        }

        err = w.Reset(&git.ResetOptions{
            Mode:   git.HardReset,
            Commit: remoteRef.Hash(),
        })
        if err != nil {
            if !silent {
                fmt.Printf("[err] failed to reset repository %s to %s: %v\n", repoPath, remoteBranchRef, err)
            }
            continue
        }

        if !silent {
            fmt.Printf("[inf] repository %s updated successfully to %s.\n", repoPath, defaultBranch)
        }
    }
}

func extractRepoNameFromURL(repoURL string) string {
    parts := strings.Split(repoURL, "/")
    return parts[len(parts)-2] + "_" + parts[len(parts)-1]
}