etchosts.go 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208
  1. package etchosts
  2. import (
  3. "bufio"
  4. "bytes"
  5. "fmt"
  6. "io"
  7. "io/ioutil"
  8. "os"
  9. "regexp"
  10. "strings"
  11. "sync"
  12. )
  13. // Record Structure for a single host record
  14. type Record struct {
  15. Hosts string
  16. IP string
  17. }
  18. // WriteTo writes record to file and returns bytes written or error
  19. func (r Record) WriteTo(w io.Writer) (int64, error) {
  20. n, err := fmt.Fprintf(w, "%s\t%s\n", r.IP, r.Hosts)
  21. return int64(n), err
  22. }
  23. var (
  24. // Default hosts config records slice
  25. defaultContent = []Record{
  26. {Hosts: "localhost", IP: "127.0.0.1"},
  27. {Hosts: "localhost ip6-localhost ip6-loopback", IP: "::1"},
  28. {Hosts: "ip6-localnet", IP: "fe00::0"},
  29. {Hosts: "ip6-mcastprefix", IP: "ff00::0"},
  30. {Hosts: "ip6-allnodes", IP: "ff02::1"},
  31. {Hosts: "ip6-allrouters", IP: "ff02::2"},
  32. }
  33. // A cache of path level locks for synchronizing /etc/hosts
  34. // updates on a file level
  35. pathMap = make(map[string]*sync.Mutex)
  36. // A package level mutex to synchronize the cache itself
  37. pathMutex sync.Mutex
  38. )
  39. func pathLock(path string) func() {
  40. pathMutex.Lock()
  41. defer pathMutex.Unlock()
  42. pl, ok := pathMap[path]
  43. if !ok {
  44. pl = &sync.Mutex{}
  45. pathMap[path] = pl
  46. }
  47. pl.Lock()
  48. return func() {
  49. pl.Unlock()
  50. }
  51. }
  52. // Drop drops the path string from the path cache
  53. func Drop(path string) {
  54. pathMutex.Lock()
  55. defer pathMutex.Unlock()
  56. delete(pathMap, path)
  57. }
  58. // Build function
  59. // path is path to host file string required
  60. // IP, hostname, and domainname set main record leave empty for no master record
  61. // extraContent is an array of extra host records.
  62. func Build(path, IP, hostname, domainname string, extraContent []Record) error {
  63. defer pathLock(path)()
  64. content := bytes.NewBuffer(nil)
  65. if IP != "" {
  66. //set main record
  67. var mainRec Record
  68. mainRec.IP = IP
  69. // User might have provided a FQDN in hostname or split it across hostname
  70. // and domainname. We want the FQDN and the bare hostname.
  71. fqdn := hostname
  72. if domainname != "" {
  73. fqdn = fmt.Sprintf("%s.%s", fqdn, domainname)
  74. }
  75. parts := strings.SplitN(fqdn, ".", 2)
  76. if len(parts) == 2 {
  77. mainRec.Hosts = fmt.Sprintf("%s %s", fqdn, parts[0])
  78. } else {
  79. mainRec.Hosts = fqdn
  80. }
  81. if _, err := mainRec.WriteTo(content); err != nil {
  82. return err
  83. }
  84. }
  85. // Write defaultContent slice to buffer
  86. for _, r := range defaultContent {
  87. if _, err := r.WriteTo(content); err != nil {
  88. return err
  89. }
  90. }
  91. // Write extra content from function arguments
  92. for _, r := range extraContent {
  93. if _, err := r.WriteTo(content); err != nil {
  94. return err
  95. }
  96. }
  97. return ioutil.WriteFile(path, content.Bytes(), 0644)
  98. }
  99. // Add adds an arbitrary number of Records to an already existing /etc/hosts file
  100. func Add(path string, recs []Record) error {
  101. defer pathLock(path)()
  102. if len(recs) == 0 {
  103. return nil
  104. }
  105. b, err := mergeRecords(path, recs)
  106. if err != nil {
  107. return err
  108. }
  109. return ioutil.WriteFile(path, b, 0644)
  110. }
  111. func mergeRecords(path string, recs []Record) ([]byte, error) {
  112. f, err := os.Open(path)
  113. if err != nil {
  114. return nil, err
  115. }
  116. defer f.Close()
  117. content := bytes.NewBuffer(nil)
  118. if _, err := content.ReadFrom(f); err != nil {
  119. return nil, err
  120. }
  121. for _, r := range recs {
  122. if _, err := r.WriteTo(content); err != nil {
  123. return nil, err
  124. }
  125. }
  126. return content.Bytes(), nil
  127. }
  128. // Delete deletes an arbitrary number of Records already existing in /etc/hosts file
  129. func Delete(path string, recs []Record) error {
  130. defer pathLock(path)()
  131. if len(recs) == 0 {
  132. return nil
  133. }
  134. old, err := os.Open(path)
  135. if err != nil {
  136. return err
  137. }
  138. var buf bytes.Buffer
  139. s := bufio.NewScanner(old)
  140. eol := []byte{'\n'}
  141. loop:
  142. for s.Scan() {
  143. b := s.Bytes()
  144. if len(b) == 0 {
  145. continue
  146. }
  147. if b[0] == '#' {
  148. buf.Write(b)
  149. buf.Write(eol)
  150. continue
  151. }
  152. for _, r := range recs {
  153. if bytes.HasSuffix(b, []byte("\t"+r.Hosts)) {
  154. continue loop
  155. }
  156. }
  157. buf.Write(b)
  158. buf.Write(eol)
  159. }
  160. old.Close()
  161. if err := s.Err(); err != nil {
  162. return err
  163. }
  164. return ioutil.WriteFile(path, buf.Bytes(), 0644)
  165. }
  166. // Update all IP addresses where hostname matches.
  167. // path is path to host file
  168. // IP is new IP address
  169. // hostname is hostname to search for to replace IP
  170. func Update(path, IP, hostname string) error {
  171. defer pathLock(path)()
  172. old, err := ioutil.ReadFile(path)
  173. if err != nil {
  174. return err
  175. }
  176. var re = regexp.MustCompile(fmt.Sprintf("(\\S*)(\\t%s)(\\s|\\.)", regexp.QuoteMeta(hostname)))
  177. return ioutil.WriteFile(path, re.ReplaceAll(old, []byte(IP+"$2"+"$3")), 0644)
  178. }