package main
import (
"flag"
"fmt"
"io"
"net/http"
"net/url"
"os"
"strings"
"github.com/PuerkitoBio/goquery"
"github.com/gdamore/tcell/v2"
"github.com/mattn/go-isatty"
"github.com/rivo/tview"
)
func newView() *tview.TextView {
return tview.NewTextView().
SetTextAlign(tview.AlignLeft).
SetText("Loading html document...")
}
func newSearchBox(onChange func(text string)) *tview.InputField {
return tview.NewInputField().
SetLabel("Enter a CSS Selector").
SetChangedFunc(onChange)
}
func quit(message string, exitCode int) {
if message != "" {
fmt.Println(message)
}
os.Exit(exitCode)
}
func fetchPage(urlStr string) (string, error) {
prefixedUrl := urlStr
if !strings.HasPrefix(prefixedUrl, "http://") && !strings.HasPrefix(prefixedUrl, "https://") {
prefixedUrl = "https://" + urlStr
}
req, err := http.NewRequest("GET", prefixedUrl, nil)
if err != nil {
return "", err
}
res, err := http.DefaultClient.Do(req)
if err != nil {
return "", err
}
defer res.Body.Close()
buf := new(strings.Builder)
_, err = io.Copy(buf, res.Body)
if err != nil {
return "", err
}
return buf.String(), nil
}
func isTerminal(fd uintptr) bool {
return isatty.IsTerminal(fd) || isatty.IsCygwinTerminal(fd)
}
func processSearch(selector string, doc *goquery.Document) (string, int) {
if len(selector) == 0 || doc == nil {
return "", 0
}
results := doc.Find(selector)
if len(results.Nodes) == 0 {
return "", 0
}
report := ""
for i := range results.Nodes {
htmlString, err := goquery.OuterHtml(results.Eq(i))
if err != nil {
report += fmt.Sprintf("error on result %d: %+v \n", i, err)
} else {
report += htmlString + "\n"
}
}
return report, len(results.Nodes)
}
func checkFileExists(str string) bool {
if _, err := os.Stat(str); err == nil {
return true
}
// if !os.ErrNotExist, still don't assume it exists
return false
}
func isUrl(str string) bool {
prefixedUrl := str
if !strings.HasPrefix(prefixedUrl, "http://") && !strings.HasPrefix(prefixedUrl, "https://") {
prefixedUrl = "https://" + str
}
u, err := url.Parse(prefixedUrl)
return err == nil && u.Scheme != "" && u.Host != ""
}
func loadDocument(urls []string) (string, error) {
source := ""
if !isTerminal(os.Stdin.Fd()) {
bytes, err := io.ReadAll(os.Stdin)
if err != nil {
return "", err
}
source = string(bytes)
} else {
if len(urls) == 0 {
return "", fmt.Errorf("Enter a file or url to search")
}
source = urls[0]
}
if len(source) < 10000 {
if isUrl(source) {
return fetchPage(source)
}
if checkFileExists(source) {
file, err := os.Open(source)
if err != nil {
return "", fmt.Errorf("Unable to open file %q: %v", source, err)
}
fileBytes, err := io.ReadAll(file)
if err != nil {
return "", fmt.Errorf("Unable to read file %q: %v", source, err)
}
return string(fileBytes), nil
}
}
if _, err := goquery.NewDocumentFromReader(strings.NewReader(source)); err == nil {
return source, nil
}
return "", fmt.Errorf("unable to load document from %q", source[:100])
}
func main() {
initialQuery := flag.String("q", "", "initial query for lmnt")
flag.Parse()
urls := flag.Args()
htmlString, err := loadDocument(urls)
if err != nil {
quit(fmt.Sprintf("Failed to load html: %+v", err), 1)
}
goqueryDoc, err := goquery.NewDocumentFromReader(strings.NewReader(htmlString))
if err != nil {
quit(fmt.Sprintf("Unable to parse document %+v", err), 1)
}
if !isTerminal(os.Stdout.Fd()) {
results, _ := processSearch(*initialQuery, goqueryDoc)
fmt.Println(results)
quit("", 0)
}
view := newView()
view.SetText(htmlString)
onSearch := func(selector string) {
report, count := processSearch(selector, goqueryDoc)
if report == "" && count == 0 {
if selector == "" {
view.SetText(htmlString)
} else {
view.SetText(fmt.Sprintf("No nodes found for %s", selector))
}
}
view.SetText(report)
}
input := newSearchBox(onSearch)
grid := tview.NewGrid().
SetColumns(0).
SetRows(1, 0).
SetBorders(true).
AddItem(input, 0, 0, 1, 1, 0, 0, true).
AddItem(view, 1, 0, 1, 1, 0, 0, false)
app := tview.NewApplication().SetRoot(grid, true)
app.SetInputCapture(func(event *tcell.EventKey) *tcell.EventKey {
if event != nil && event.Key() == tcell.KeyTab {
if view.HasFocus() {
app.SetFocus(input)
} else {
app.SetFocus(view)
}
return nil
}
return event
})
view.SetChangedFunc(func() {
app.Draw()
})
if err := app.Run(); err != nil {
panic(err)
}
}