Last active
April 5, 2024 13:45
-
-
Save jostmart/67b2ad7507d0f175cc5768e4027270a5 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
package main | |
import ( | |
"fmt" | |
"io/ioutil" | |
"net/http" | |
"sync" | |
"time" | |
"github.com/go-yaml/yaml" | |
) | |
// Endpoint struct to hold IP address and port | |
type Endpoint struct { | |
IP string `yaml:"ip"` | |
Port string `yaml:"port"` | |
} | |
// Config struct to hold list of endpoints | |
type Config struct { | |
Endpoints []Endpoint `yaml:"endpoints"` | |
} | |
// ManagedRule struct to hold information about a managed iptables rule | |
type ManagedRule struct { | |
IP string | |
Port string | |
Added bool | |
} | |
var ( | |
rules []ManagedRule | |
rulesLock sync.Mutex | |
) | |
func main() { | |
defer func() { | |
if r := recover(); r != nil { | |
fmt.Println("Recovered from panic:", r) | |
} | |
}() | |
// Load configuration from YAML file | |
config, err := readConfig("config.yml") | |
if err != nil { | |
fmt.Println("Error reading config file:", err) | |
return | |
} | |
// Initialize rules slice | |
rules = make([]ManagedRule, 0) | |
// Loop through all endpoints | |
for _, endpoint := range config.Endpoints { | |
// Perform health check and manage iptables rule | |
err := checkAndManage(endpoint) | |
if err != nil { | |
fmt.Printf("Error checking endpoint %s:%s: %v\n", endpoint.IP, endpoint.Port, err) | |
} | |
} | |
// Keep the program running indefinitely | |
select {} | |
} | |
// Read configuration from YAML file | |
func readConfig(filename string) (Config, error) { | |
var config Config | |
data, err := ioutil.ReadFile(filename) | |
if err != nil { | |
return config, err | |
} | |
err = yaml.Unmarshal(data, &config) | |
if err != nil { | |
return config, err | |
} | |
return config, nil | |
} | |
// Perform health check and manage iptables rule | |
func checkAndManage(endpoint Endpoint) error { | |
// Perform HTTP GET request for health check | |
client := http.Client{ | |
Timeout: 5 * time.Second, | |
} | |
resp, err := client.Get(fmt.Sprintf("http://%s:%s/health", endpoint.IP, endpoint.Port)) | |
if err != nil { | |
// If health check fails, remove iptables rule if it exists | |
removeIptablesRule(endpoint.IP, endpoint.Port) | |
return fmt.Errorf("health check failed for endpoint %s:%s: %v", endpoint.IP, endpoint.Port, err) | |
} | |
defer resp.Body.Close() | |
// If status code is OK, manage iptables rule | |
if resp.StatusCode == http.StatusOK { | |
addOrUpdateIptablesRule(endpoint.IP, endpoint.Port) | |
return nil | |
} | |
// If status code is not OK, remove iptables rule if it exists | |
removeIptablesRule(endpoint.IP, endpoint.Port) | |
return fmt.Errorf("endpoint %s:%s returned status code %d", endpoint.IP, endpoint.Port, resp.StatusCode) | |
} | |
// Add or update iptables rule | |
func addOrUpdateIptablesRule(ip, port string) { | |
rulesLock.Lock() | |
defer rulesLock.Unlock() | |
// Check if rule already exists | |
for i, rule := range rules { | |
if rule.IP == ip && rule.Port == port { | |
if !rule.Added { | |
// If rule exists but not added, add it | |
addIptablesRule(ip, port) | |
rules[i].Added = true | |
} | |
return | |
} | |
} | |
// If rule doesn't exist, add it | |
addIptablesRule(ip, port) | |
rules = append(rules, ManagedRule{IP: ip, Port: port, Added: true}) | |
} | |
// Remove iptables rule | |
func removeIptablesRule(ip, port string) { | |
rulesLock.Lock() | |
defer rulesLock.Unlock() | |
// Check if rule exists | |
for i, rule := range rules { | |
if rule.IP == ip && rule.Port == port && rule.Added { | |
// If rule exists and added, remove it | |
removeIptablesRuleExec(ip, port) | |
rules[i].Added = false | |
return | |
} | |
} | |
} | |
// Add iptables rule using exec | |
func addIptablesRule(ip, port string) { | |
// Simulated add iptables rule | |
fmt.Printf("Adding iptables rule for %s:%s\n", ip, port) | |
} | |
// Remove iptables rule using exec | |
func removeIptablesRuleExec(ip, port string) { | |
// Simulated remove iptables rule | |
fmt.Printf("Removing iptables rule for %s:%s\n", ip, port) | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment