Simple hosts file management in Go (golang).
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

259 lines
5.0 KiB

  1. package goodhosts
  2. import (
  3. "bufio"
  4. "errors"
  5. "fmt"
  6. "net"
  7. "os"
  8. )
  9. const sectionStart = "### Start Hosts for"
  10. const sectionEnd = "### End Hosts for"
  11. // Hosts Represents a hosts file.
  12. type Hosts struct {
  13. Path string
  14. Section string
  15. FileLines []HostsLine
  16. SectionLines []HostsLine
  17. }
  18. // IsWritable Return ```true``` if hosts file is writable.
  19. func (h *Hosts) IsWritable() bool {
  20. _, err := os.OpenFile(h.Path, os.O_WRONLY, 0660)
  21. if err != nil {
  22. return false
  23. }
  24. return true
  25. }
  26. // Load the hosts file into ```l.Lines```.
  27. // ```Load()``` is called by ```NewHosts()``` and ```Hosts.Flush()``` so you
  28. // generally you won't need to call this yourself.
  29. func (h *Hosts) Load() error {
  30. var fileLines []HostsLine
  31. var sectionLines []HostsLine
  32. var inSection bool
  33. file, err := os.Open(h.Path)
  34. if err != nil {
  35. return err
  36. }
  37. defer file.Close()
  38. scanner := bufio.NewScanner(file)
  39. for scanner.Scan() {
  40. line := NewHostsLine(scanner.Text())
  41. if err != nil {
  42. return err
  43. }
  44. if line.Raw == fmt.Sprintf("%s %s", sectionEnd, h.Section) {
  45. inSection = false
  46. }
  47. if inSection {
  48. sectionLines = append(sectionLines, line)
  49. }
  50. if line.Raw == fmt.Sprintf("%s %s", sectionStart, h.Section) {
  51. inSection = true
  52. }
  53. if !inSection && line.Raw != fmt.Sprintf("%s %s", sectionEnd, h.Section) {
  54. fileLines = append(fileLines, line)
  55. }
  56. }
  57. if err := scanner.Err(); err != nil {
  58. return err
  59. }
  60. h.FileLines = fileLines
  61. h.SectionLines = sectionLines
  62. return nil
  63. }
  64. // Flush any changes made to hosts file.
  65. func (h Hosts) Flush() error {
  66. file, err := os.Create(h.Path)
  67. if err != nil {
  68. return err
  69. }
  70. if len(h.SectionLines) > 0 {
  71. if len(h.Section) > 0 {
  72. h.FileLines = append(h.FileLines, NewHostsLine(""))
  73. h.FileLines = append(h.FileLines, NewHostsLine(fmt.Sprintf("%s %s", sectionStart, h.Section)))
  74. }
  75. for _, sectionLine := range h.SectionLines {
  76. h.FileLines = append(h.FileLines, sectionLine)
  77. }
  78. if len(h.Section) > 0 {
  79. h.FileLines = append(h.FileLines, NewHostsLine(fmt.Sprintf("%s %s", sectionEnd, h.Section)))
  80. h.FileLines = append(h.FileLines, NewHostsLine(""))
  81. }
  82. }
  83. var isBlank bool
  84. w := bufio.NewWriter(file)
  85. for _, line := range h.FileLines {
  86. if !isBlank || len(line.Raw) > 1 {
  87. fmt.Fprintf(w, "%s%s", line.Raw, eol)
  88. }
  89. if len(line.Raw) < 2 {
  90. isBlank = true
  91. } else {
  92. isBlank = false
  93. }
  94. }
  95. err = w.Flush()
  96. if err != nil {
  97. return err
  98. }
  99. return h.Load()
  100. }
  101. // Add an entry to the hosts file.
  102. func (h *Hosts) Add(ip, comment string, hosts ...string) error {
  103. if net.ParseIP(ip) == nil {
  104. return fmt.Errorf("%q is an invalid IP address", ip)
  105. }
  106. for _, host := range hosts {
  107. if h.Has(ip, host, true) {
  108. return fmt.Errorf("%s has already been assigned", host)
  109. }
  110. if !h.Has(ip, host, false) {
  111. endLine := NewHostsLine(buildRawLine(ip, host, comment))
  112. endLine.Comment = comment
  113. h.SectionLines = append(h.SectionLines, endLine)
  114. }
  115. }
  116. return nil
  117. }
  118. // Has Return a bool if ip/host combo in hosts file.
  119. func (h *Hosts) Has(ip string, host string, forceFile bool) bool {
  120. pos := h.getHostPosition(ip, host, forceFile)
  121. return pos != -1
  122. }
  123. // RemoveSection removes an entire section from the hostsfile
  124. func (h *Hosts) RemoveSection() error {
  125. if len(h.Section) == 0 {
  126. return errors.New("No section Provided")
  127. }
  128. h.SectionLines = nil
  129. return nil
  130. }
  131. // Remove an entry from the hosts file.
  132. func (h *Hosts) Remove(ip string, hosts ...string) error {
  133. var outputLines []HostsLine
  134. inputLines := h.SectionLines
  135. if len(h.Section) == 0 {
  136. inputLines = h.FileLines
  137. }
  138. if net.ParseIP(ip) == nil {
  139. return fmt.Errorf("%q is an invalid IP address", ip)
  140. }
  141. for _, line := range inputLines {
  142. // Bad lines or comments just get readded.
  143. if line.Err != nil || IsComment(line.Raw) || line.IP != ip {
  144. outputLines = append(outputLines, line)
  145. continue
  146. }
  147. var newHosts []string
  148. for _, checkHost := range line.Hosts {
  149. if !itemInSlice(checkHost, hosts) {
  150. newHosts = append(newHosts, checkHost)
  151. }
  152. }
  153. // If hosts is empty, skip the line completely.
  154. if len(newHosts) > 0 {
  155. newLineRaw := line.IP
  156. for _, host := range newHosts {
  157. newLineRaw = fmt.Sprintf("%s %s", newLineRaw, host)
  158. }
  159. if len(line.Comment) > 0 {
  160. newLineRaw = fmt.Sprintf("%s #%s", newLineRaw, line.Comment)
  161. }
  162. newLine := NewHostsLine(newLineRaw)
  163. outputLines = append(outputLines, newLine)
  164. }
  165. }
  166. if len(h.Section) == 0 {
  167. h.FileLines = outputLines
  168. } else {
  169. h.SectionLines = outputLines
  170. }
  171. return nil
  172. }
  173. func (h Hosts) getHostPosition(ip string, host string, forceFile bool) int {
  174. checkLines := h.FileLines
  175. if len(h.Section) > 0 && !forceFile {
  176. checkLines = h.SectionLines
  177. }
  178. for i := range checkLines {
  179. line := checkLines[i]
  180. if !IsComment(line.Raw) && line.Raw != "" {
  181. if ip == line.IP && itemInSlice(host, line.Hosts) {
  182. return i
  183. }
  184. }
  185. }
  186. return -1
  187. }
  188. func (h Hosts) getIPPosition(ip string) int {
  189. for i := range h.FileLines {
  190. line := h.FileLines[i]
  191. if !IsComment(line.Raw) && line.Raw != "" {
  192. if line.IP == ip {
  193. return i
  194. }
  195. }
  196. }
  197. return -1
  198. }