├── .gitignore ├── LICENSE ├── README.md ├── cmd ├── copy.go ├── root.go └── run.go ├── hosts ├── config.go └── files.go ├── main.go └── sessions ├── auth.go ├── sftp.go └── ssh.go /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | 3 | glide.* 4 | vendor/ 5 | 6 | hosts.list 7 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | 2 | Apache License 3 | Version 2.0, January 2004 4 | http://www.apache.org/licenses/ 5 | 6 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 7 | 8 | 1. Definitions. 9 | 10 | "License" shall mean the terms and conditions for use, reproduction, 11 | and distribution as defined by Sections 1 through 9 of this document. 12 | 13 | "Licensor" shall mean the copyright owner or entity authorized by 14 | the copyright owner that is granting the License. 15 | 16 | "Legal Entity" shall mean the union of the acting entity and all 17 | other entities that control, are controlled by, or are under common 18 | control with that entity. For the purposes of this definition, 19 | "control" means (i) the power, direct or indirect, to cause the 20 | direction or management of such entity, whether by contract or 21 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 22 | outstanding shares, or (iii) beneficial ownership of such entity. 23 | 24 | "You" (or "Your") shall mean an individual or Legal Entity 25 | exercising permissions granted by this License. 26 | 27 | "Source" form shall mean the preferred form for making modifications, 28 | including but not limited to software source code, documentation 29 | source, and configuration files. 30 | 31 | "Object" form shall mean any form resulting from mechanical 32 | transformation or translation of a Source form, including but 33 | not limited to compiled object code, generated documentation, 34 | and conversions to other media types. 35 | 36 | "Work" shall mean the work of authorship, whether in Source or 37 | Object form, made available under the License, as indicated by a 38 | copyright notice that is included in or attached to the work 39 | (an example is provided in the Appendix below). 40 | 41 | "Derivative Works" shall mean any work, whether in Source or Object 42 | form, that is based on (or derived from) the Work and for which the 43 | editorial revisions, annotations, elaborations, or other modifications 44 | represent, as a whole, an original work of authorship. For the purposes 45 | of this License, Derivative Works shall not include works that remain 46 | separable from, or merely link (or bind by name) to the interfaces of, 47 | the Work and Derivative Works thereof. 48 | 49 | "Contribution" shall mean any work of authorship, including 50 | the original version of the Work and any modifications or additions 51 | to that Work or Derivative Works thereof, that is intentionally 52 | submitted to Licensor for inclusion in the Work by the copyright owner 53 | or by an individual or Legal Entity authorized to submit on behalf of 54 | the copyright owner. For the purposes of this definition, "submitted" 55 | means any form of electronic, verbal, or written communication sent 56 | to the Licensor or its representatives, including but not limited to 57 | communication on electronic mailing lists, source code control systems, 58 | and issue tracking systems that are managed by, or on behalf of, the 59 | Licensor for the purpose of discussing and improving the Work, but 60 | excluding communication that is conspicuously marked or otherwise 61 | designated in writing by the copyright owner as "Not a Contribution." 62 | 63 | "Contributor" shall mean Licensor and any individual or Legal Entity 64 | on behalf of whom a Contribution has been received by Licensor and 65 | subsequently incorporated within the Work. 66 | 67 | 2. Grant of Copyright License. Subject to the terms and conditions of 68 | this License, each Contributor hereby grants to You a perpetual, 69 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 70 | copyright license to reproduce, prepare Derivative Works of, 71 | publicly display, publicly perform, sublicense, and distribute the 72 | Work and such Derivative Works in Source or Object form. 73 | 74 | 3. Grant of Patent License. Subject to the terms and conditions of 75 | this License, each Contributor hereby grants to You a perpetual, 76 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 77 | (except as stated in this section) patent license to make, have made, 78 | use, offer to sell, sell, import, and otherwise transfer the Work, 79 | where such license applies only to those patent claims licensable 80 | by such Contributor that are necessarily infringed by their 81 | Contribution(s) alone or by combination of their Contribution(s) 82 | with the Work to which such Contribution(s) was submitted. If You 83 | institute patent litigation against any entity (including a 84 | cross-claim or counterclaim in a lawsuit) alleging that the Work 85 | or a Contribution incorporated within the Work constitutes direct 86 | or contributory patent infringement, then any patent licenses 87 | granted to You under this License for that Work shall terminate 88 | as of the date such litigation is filed. 89 | 90 | 4. Redistribution. You may reproduce and distribute copies of the 91 | Work or Derivative Works thereof in any medium, with or without 92 | modifications, and in Source or Object form, provided that You 93 | meet the following conditions: 94 | 95 | (a) You must give any other recipients of the Work or 96 | Derivative Works a copy of this License; and 97 | 98 | (b) You must cause any modified files to carry prominent notices 99 | stating that You changed the files; and 100 | 101 | (c) You must retain, in the Source form of any Derivative Works 102 | that You distribute, all copyright, patent, trademark, and 103 | attribution notices from the Source form of the Work, 104 | excluding those notices that do not pertain to any part of 105 | the Derivative Works; and 106 | 107 | (d) If the Work includes a "NOTICE" text file as part of its 108 | distribution, then any Derivative Works that You distribute must 109 | include a readable copy of the attribution notices contained 110 | within such NOTICE file, excluding those notices that do not 111 | pertain to any part of the Derivative Works, in at least one 112 | of the following places: within a NOTICE text file distributed 113 | as part of the Derivative Works; within the Source form or 114 | documentation, if provided along with the Derivative Works; or, 115 | within a display generated by the Derivative Works, if and 116 | wherever such third-party notices normally appear. The contents 117 | of the NOTICE file are for informational purposes only and 118 | do not modify the License. You may add Your own attribution 119 | notices within Derivative Works that You distribute, alongside 120 | or as an addendum to the NOTICE text from the Work, provided 121 | that such additional attribution notices cannot be construed 122 | as modifying the License. 123 | 124 | You may add Your own copyright statement to Your modifications and 125 | may provide additional or different license terms and conditions 126 | for use, reproduction, or distribution of Your modifications, or 127 | for any such Derivative Works as a whole, provided Your use, 128 | reproduction, and distribution of the Work otherwise complies with 129 | the conditions stated in this License. 130 | 131 | 5. Submission of Contributions. Unless You explicitly state otherwise, 132 | any Contribution intentionally submitted for inclusion in the Work 133 | by You to the Licensor shall be under the terms and conditions of 134 | this License, without any additional terms or conditions. 135 | Notwithstanding the above, nothing herein shall supersede or modify 136 | the terms of any separate license agreement you may have executed 137 | with Licensor regarding such Contributions. 138 | 139 | 6. Trademarks. This License does not grant permission to use the trade 140 | names, trademarks, service marks, or product names of the Licensor, 141 | except as required for reasonable and customary use in describing the 142 | origin of the Work and reproducing the content of the NOTICE file. 143 | 144 | 7. Disclaimer of Warranty. Unless required by applicable law or 145 | agreed to in writing, Licensor provides the Work (and each 146 | Contributor provides its Contributions) on an "AS IS" BASIS, 147 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 148 | implied, including, without limitation, any warranties or conditions 149 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 150 | PARTICULAR PURPOSE. You are solely responsible for determining the 151 | appropriateness of using or redistributing the Work and assume any 152 | risks associated with Your exercise of permissions under this License. 153 | 154 | 8. Limitation of Liability. In no event and under no legal theory, 155 | whether in tort (including negligence), contract, or otherwise, 156 | unless required by applicable law (such as deliberate and grossly 157 | negligent acts) or agreed to in writing, shall any Contributor be 158 | liable to You for damages, including any direct, indirect, special, 159 | incidental, or consequential damages of any character arising as a 160 | result of this License or out of the use or inability to use the 161 | Work (including but not limited to damages for loss of goodwill, 162 | work stoppage, computer failure or malfunction, or any and all 163 | other commercial damages or losses), even if such Contributor 164 | has been advised of the possibility of such damages. 165 | 166 | 9. Accepting Warranty or Additional Liability. While redistributing 167 | the Work or Derivative Works thereof, You may choose to offer, 168 | and charge a fee for, acceptance of support, warranty, indemnity, 169 | or other liability obligations and/or rights consistent with this 170 | License. However, in accepting such obligations, You may act only 171 | on Your own behalf and on Your sole responsibility, not on behalf 172 | of any other Contributor, and only if You agree to indemnify, 173 | defend, and hold each Contributor harmless for any liability 174 | incurred by, or claims asserted against, such Contributor by reason 175 | of your accepting any such warranty or additional liability. 176 | 177 | END OF TERMS AND CONDITIONS 178 | 179 | APPENDIX: How to apply the Apache License to your work. 180 | 181 | To apply the Apache License to your work, attach the following 182 | boilerplate notice, with the fields enclosed by brackets "[]" 183 | replaced with your own identifying information. (Don't include 184 | the brackets!) The text should be enclosed in the appropriate 185 | comment syntax for the file format. We also recommend that a 186 | file or class name and description of purpose be included on the 187 | same "printed page" as the copyright notice for easier 188 | identification within third-party archives. 189 | 190 | Copyright [yyyy] [name of copyright owner] 191 | 192 | Licensed under the Apache License, Version 2.0 (the "License"); 193 | you may not use this file except in compliance with the License. 194 | You may obtain a copy of the License at 195 | 196 | http://www.apache.org/licenses/LICENSE-2.0 197 | 198 | Unless required by applicable law or agreed to in writing, software 199 | distributed under the License is distributed on an "AS IS" BASIS, 200 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 201 | See the License for the specific language governing permissions and 202 | limitations under the License. 203 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Omnitool 2 | 3 | Omnitool is a tool for using SSH on multiple machines in parallel. It is particularly useful for times when you need a bunch of machines for a short period of time but don't want to invest the time in building up a more robust management system. 4 | 5 | Omnitool's goal is to let you think in terms of one machine while working with N machines. 6 | 7 | ## Installing 8 | 9 | ```shell 10 | $ go get github.com/jmsdnns/omnitool 11 | $ go install github.com/jmsdnns/omnitool 12 | ``` 13 | 14 | ## Using It 15 | 16 | Omnitool has help built-in. 17 | 18 | ```shell 19 | Usage: 20 | omnitool [command] 21 | 22 | Available Commands: 23 | copy Copies file to host group 24 | run Runs a command on host group 25 | 26 | Flags: 27 | -g, --group string host group for task 28 | -h, --help help for omnitool 29 | --hostsfile string path to hosts file (default "hosts.list") 30 | -k, --keyfile string path to ssh key 31 | -u, --username string username for ssh 32 | 33 | Use "omnitool [command] --help" for more information about a command. 34 | ``` 35 | 36 | Running a command on a host group looks like this: 37 | 38 | ```shell 39 | $ ./omnitool -u ... -k ... -g ... run "ls -l" 40 | CMD: ls -l 41 | 42 | Host: 33.33.33.12:22 43 | Result: 44 | total 0 45 | -rw-rw-r-- 1 vagrant vagrant 0 Aug 14 02:53 machine_2 46 | 47 | Host: 33.33.33.11:22 48 | Result: 49 | total 0 50 | -rw-rw-r-- 1 vagrant vagrant 0 Aug 14 02:53 machine_1 51 | ``` 52 | 53 | Copying a file to a host group looks like this: 54 | 55 | ```shell 56 | $ ./omnitool -u ... -k ... -g ... copy hosts.list hosts.list 57 | Host: 33.33.33.12:22 58 | Result: ok 59 | 60 | Host: 33.33.33.11:22 61 | Result: ok 62 | ``` 63 | 64 | ## Hosts File 65 | 66 | A hosts file looks like this: 67 | 68 | ``` 69 | [vagrants] 70 | 127.0.0.1:2222 71 | 127.0.0.1:2200 72 | 73 | [jms labs] 74 | 192.168.0.5:22 75 | 192.168.0.6:22 76 | ``` 77 | 78 | ## Using with Vagrant 79 | 80 | Using Omnitool with Vagrant is easy, but you will need to tell Vagrant not to automatically generate SSH keys for each host. This lets us use a single key for each host, which more accurately reflects what you'd have in AWS anyway. 81 | 82 | _For more info, read about the `config.ssh.insert_key` flag in [Vagrant's SSH documentation](https://www.vagrantup.com/docs/vagrantfile/ssh_settings.html)._ 83 | 84 | | Flag | Default value | 85 | | ------ | ------------------------------------- | 86 | | `-u` | vagrant | 87 | | `-k` | $HOME/.vagrant.d/insecure_private_key | 88 | -------------------------------------------------------------------------------- /cmd/copy.go: -------------------------------------------------------------------------------- 1 | package cmd 2 | 3 | import ( 4 | "fmt" 5 | "log" 6 | "time" 7 | 8 | "github.com/jmsdnns/omnitool/sessions" 9 | "github.com/spf13/cobra" 10 | ) 11 | 12 | // sftpCmd represents the sftp command 13 | var sftpCmd = &cobra.Command{ 14 | Use: "copy", 15 | Short: "Copies file to host group", 16 | Long: `Creates an SFTP pool to a host group and copies file to each host. Success of 17 | failure is then collected from each host and reported here`, 18 | Run: cmdCopy, 19 | } 20 | 21 | func init() { 22 | RootCmd.AddCommand(sftpCmd) 23 | } 24 | 25 | func cmdCopy(cmd *cobra.Command, args []string) { 26 | username := cmd.Flags().Lookup("username").Value.String() 27 | keyfile := cmd.Flags().Lookup("keyfile").Value.String() 28 | hostsfile := cmd.Flags().Lookup("hostsfile").Value.String() 29 | group := cmd.Flags().Lookup("group").Value.String() 30 | 31 | hostList, err := ParseHostArgs(hostsfile, group) 32 | if err != nil { 33 | log.Fatal(err) 34 | return 35 | } 36 | 37 | if len(args) != 2 { 38 | log.Fatal("localPath and remotePath arguments not found") 39 | return 40 | } 41 | localPath := args[0] 42 | remotePath := args[1] 43 | 44 | results := make(chan sessions.SFTPResponse) 45 | timeout := time.After(60 * time.Second) 46 | sessions.MapCopy(hostList, username, keyfile, localPath, remotePath, results) 47 | 48 | for i := 0; i < len(hostList); i++ { 49 | select { 50 | case r := <-results: 51 | fmt.Printf("Host: %s\n", r.Host) 52 | fmt.Printf("Result: %s\n\n", r.Result) 53 | case <-timeout: 54 | fmt.Println("Timed out!") 55 | } 56 | } 57 | } 58 | -------------------------------------------------------------------------------- /cmd/root.go: -------------------------------------------------------------------------------- 1 | package cmd 2 | 3 | import ( 4 | "fmt" 5 | "os" 6 | 7 | "github.com/jmsdnns/omnitool/hosts" 8 | "github.com/spf13/cobra" 9 | ) 10 | 11 | // RootCmd represents the base command when called without any subcommands 12 | var RootCmd = &cobra.Command{ 13 | Use: "omnitool", 14 | Short: "A tool for managing machines via parallel SSH pools", 15 | Long: `Omnitool's goal is to let you think in terms of one machine while working with 16 | N machines.`, 17 | } 18 | 19 | // Execute adds all child commands to the root command sets flags appropriately. 20 | // This is called by main.main(). It only needs to happen once to the rootCmd. 21 | func Execute() { 22 | if err := RootCmd.Execute(); err != nil { 23 | fmt.Println(err) 24 | os.Exit(-1) 25 | } 26 | } 27 | 28 | func init() { 29 | RootCmd.PersistentFlags().StringP("username", "u", "", "username for ssh") 30 | RootCmd.PersistentFlags().StringP("keyfile", "k", "", "path to ssh key") 31 | RootCmd.PersistentFlags().StringP("hostsfile", "", "hosts.list", "path to hosts file") 32 | RootCmd.PersistentFlags().StringP("group", "g", "", "host group for task") 33 | } 34 | 35 | // ParseHostArgs reads a hosts file, finds the requested group, and returns 36 | // the corresponding list of hosts 37 | func ParseHostArgs(hostsfile string, group string) ([]string, error) { 38 | // Load machine list file 39 | hostsConfig, err := hosts.LoadHostsFile(hostsfile) 40 | if err != nil { 41 | return nil, err 42 | } 43 | 44 | hostGroup := hostsConfig.Get(group) 45 | return hostGroup, nil 46 | } 47 | -------------------------------------------------------------------------------- /cmd/run.go: -------------------------------------------------------------------------------- 1 | package cmd 2 | 3 | import ( 4 | "fmt" 5 | "log" 6 | "strings" 7 | "time" 8 | 9 | "github.com/jmsdnns/omnitool/sessions" 10 | "github.com/spf13/cobra" 11 | ) 12 | 13 | // runCmd represents the run command 14 | var runCmd = &cobra.Command{ 15 | Use: "run", 16 | Short: "Runs a command on host group", 17 | Long: `Creates an SSH pool to a host group and executes a command. The output is then 18 | collected from each host and displayed here.`, 19 | Run: cmdRun, 20 | } 21 | 22 | func init() { 23 | RootCmd.AddCommand(runCmd) 24 | } 25 | 26 | func cmdRun(cmd *cobra.Command, args []string) { 27 | username := cmd.Flags().Lookup("username").Value.String() 28 | keyfile := cmd.Flags().Lookup("keyfile").Value.String() 29 | hostsfile := cmd.Flags().Lookup("hostsfile").Value.String() 30 | group := cmd.Flags().Lookup("group").Value.String() 31 | 32 | hostList, err := ParseHostArgs(hostsfile, group) 33 | if err != nil { 34 | log.Fatal(err) 35 | return 36 | } 37 | 38 | userCmd := strings.Join(args, " ") 39 | 40 | results := make(chan sessions.SSHResponse) 41 | timeout := time.After(60 * time.Second) 42 | sessions.MapCmd(hostList, username, keyfile, userCmd, results) 43 | 44 | fmt.Printf("CMD: %s\n\n", userCmd) 45 | for i := 0; i < len(hostList); i++ { 46 | select { 47 | case r := <-results: 48 | fmt.Printf("Host: %s\n", r.Host) 49 | fmt.Printf("Result:\n%s\n", r.Result) 50 | case <-timeout: 51 | fmt.Println("Timed out!") 52 | } 53 | } 54 | 55 | } 56 | -------------------------------------------------------------------------------- /hosts/config.go: -------------------------------------------------------------------------------- 1 | package hosts 2 | 3 | // This an interface for reading ini-like host groups from file. The format is 4 | // simply a group name with hostnames on each line below, until a new group 5 | // starts or EOF. 6 | // 7 | // Example: 8 | // 9 | // [apiservers] 10 | // 10.0.0.1 11 | // 10.0.0.2 12 | // 10.0.0.3 13 | // 10.0.0.4 14 | // 15 | // [dbservers] 16 | // 10.0.0.5 17 | // 10.0.0.6 18 | // 19 | // This file is stored as a map[group name] => list of hosts 20 | 21 | import ( 22 | "bufio" 23 | "os" 24 | ) 25 | 26 | // Config is a map of group name => group 27 | type Config map[string]Group 28 | 29 | // Group is a list of hosts 30 | type Group []string 31 | 32 | // Get takes a group name and returns a HostGroup 33 | func (c Config) Get(groupName string) Group { 34 | group := c[groupName] 35 | if group == nil { 36 | group = make(Group, 0) 37 | c[groupName] = group 38 | } 39 | 40 | return group 41 | } 42 | 43 | // Set associates a HostGroup with a group name 44 | func (c Config) Set(groupName string, group Group) { 45 | c[groupName] = group 46 | } 47 | 48 | // LoadFile opens a file by name and passes io.Reader to File.Load 49 | func (c Config) LoadFile(file string) (err error) { 50 | in, err := os.Open(file) 51 | if err != nil { 52 | return 53 | } 54 | defer in.Close() 55 | 56 | bufin := bufio.NewReader(in) 57 | 58 | err = parseFile(bufin, &c) 59 | return err 60 | } 61 | -------------------------------------------------------------------------------- /hosts/files.go: -------------------------------------------------------------------------------- 1 | package hosts 2 | 3 | import ( 4 | "bufio" 5 | "io" 6 | "regexp" 7 | "strings" 8 | ) 9 | 10 | // 11 | // File Parsing 12 | // 13 | 14 | // LoadHostsFile reads a file into a File instance 15 | func LoadHostsFile(filename string) (Config, error) { 16 | config := make(Config) 17 | err := config.LoadFile(filename) 18 | return config, err 19 | } 20 | 21 | func parseFile(in *bufio.Reader, c *Config) (err error) { 22 | // [group name] 23 | groupNamePattern := regexp.MustCompile(`^\[(.*)\]$`) 24 | 25 | // Tracks last group name found 26 | currentGroup := "" 27 | 28 | // Walk across each line input 29 | for done := false; !done; { 30 | var line string 31 | if line, err = in.ReadString('\n'); err != nil { 32 | if err == io.EOF { 33 | done = true 34 | continue 35 | } else { 36 | return err 37 | } 38 | } 39 | 40 | line = strings.TrimSpace(line) 41 | 42 | // Skip blank lines 43 | if len(line) == 0 { 44 | continue 45 | } 46 | 47 | // Skip comments 48 | if line[0] == ';' || line[0] == '#' { 49 | continue 50 | } 51 | 52 | // Group names are the only pattern we care about in the file. Anything that 53 | // isn't a group name should be treated like a host address 54 | if groups := groupNamePattern.FindStringSubmatch(line); groups != nil { 55 | name := strings.TrimSpace(groups[1]) 56 | currentGroup = name 57 | c.Get(name) 58 | } else { 59 | host := line 60 | group := c.Get(currentGroup) 61 | group = append(group, host) 62 | c.Set(currentGroup, group) 63 | } 64 | } 65 | 66 | return nil 67 | } 68 | -------------------------------------------------------------------------------- /main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import "github.com/jmsdnns/omnitool/cmd" 4 | 5 | func main() { 6 | cmd.Execute() 7 | } 8 | -------------------------------------------------------------------------------- /sessions/auth.go: -------------------------------------------------------------------------------- 1 | package sessions 2 | 3 | import ( 4 | "io/ioutil" 5 | "log" 6 | 7 | "golang.org/x/crypto/ssh" 8 | ) 9 | 10 | // 11 | // Connection Setup 12 | // 13 | 14 | func loadPrivateKey(filepath string) (ssh.Signer, error) { 15 | pemBytes, err := ioutil.ReadFile(filepath) 16 | if err != nil { 17 | log.Fatal(err) 18 | return nil, err 19 | } 20 | 21 | signer, err := ssh.ParsePrivateKey(pemBytes) 22 | if err != nil { 23 | log.Fatal(err) 24 | return nil, err 25 | } 26 | 27 | return signer, nil 28 | } 29 | 30 | func generateConfig(username string, keypath string) (*ssh.ClientConfig, error) { 31 | signer, err := loadPrivateKey(keypath) 32 | if err != nil { 33 | return nil, err 34 | } 35 | 36 | config := &ssh.ClientConfig{ 37 | User: username, 38 | Auth: []ssh.AuthMethod{ssh.PublicKeys(signer)}, 39 | } 40 | 41 | return config, nil 42 | } 43 | -------------------------------------------------------------------------------- /sessions/sftp.go: -------------------------------------------------------------------------------- 1 | package sessions 2 | 3 | import ( 4 | "fmt" 5 | "io" 6 | "log" 7 | "os" 8 | 9 | "github.com/jmsdnns/omnitool/hosts" 10 | "github.com/pkg/sftp" 11 | ) 12 | 13 | // SFTPConnnection is a container for the pieces necessary to hold an SFTP connection 14 | // open in a goroutine 15 | type SFTPConnnection struct { 16 | ssh *SSHConnection 17 | client *sftp.Client 18 | error error 19 | } 20 | 21 | func (s *SFTPConnnection) init(ss *SSHConnection) error { 22 | s.ssh = ss 23 | 24 | client, err := sftp.NewClient(ss.client) 25 | if err != nil { 26 | log.Fatal(err) 27 | return err 28 | } 29 | s.client = client 30 | 31 | return nil 32 | } 33 | 34 | // TearDown closes all of the connections in an SFTP connection 35 | func (s *SFTPConnnection) TearDown() { 36 | s.client.Close() 37 | } 38 | 39 | func (s *SFTPConnnection) copyFile(localPath string, remotePath string) error { 40 | // Create remote file handle 41 | f, err := s.client.Create(remotePath) 42 | if err != nil { 43 | fmt.Println(err.Error()) 44 | return err 45 | } 46 | 47 | // Open input file 48 | fi, err := os.Open(localPath) 49 | if err != nil { 50 | log.Fatal(err) 51 | return err 52 | } 53 | defer fi.Close() 54 | 55 | fileChunk := make([]byte, 1024) 56 | for { 57 | // read a chunk 58 | n, err := fi.Read(fileChunk) 59 | if err != nil && err != io.EOF { 60 | log.Fatal(err) 61 | } 62 | if n == 0 { 63 | break 64 | } 65 | 66 | // Write data to file handle 67 | if _, err := f.Write(fileChunk); err != nil { 68 | fmt.Println(err.Error()) 69 | return err 70 | } 71 | } 72 | 73 | return nil 74 | } 75 | 76 | // InitSFTPConnection returns an SFTP client from an existing SSHSession 77 | func InitSFTPConnection(s *SSHConnection) (*SFTPConnnection, error) { 78 | sftpConn := &SFTPConnnection{} 79 | 80 | err := sftpConn.init(s) 81 | if err != nil { 82 | log.Fatal(err) 83 | return nil, err 84 | } 85 | 86 | return sftpConn, nil 87 | } 88 | 89 | // SFTPResponse contains the result of copying a file to a host via SFTP 90 | type SFTPResponse struct { 91 | Host string 92 | Result string 93 | Err error 94 | } 95 | 96 | // MapCopy takes the details for a file transfer and maps the file transfer 97 | // across a list of hosts 98 | func MapCopy(hosts hosts.Group, username string, keypath string, localPath string, remotePath string, results chan SFTPResponse) { 99 | for _, host := range hosts { 100 | go func(host string) { 101 | response := SFTPResponse{Host: host} 102 | 103 | sshConn, err := InitSSHConnection(host, username, keypath) 104 | if err != nil { 105 | log.Fatal(err) 106 | response.Err = err 107 | } 108 | sftpConn, err := InitSFTPConnection(sshConn) 109 | defer sshConn.TearDown() 110 | defer sftpConn.TearDown() 111 | 112 | err = sftpConn.copyFile(localPath, remotePath) 113 | response.Result = "ok" 114 | if err != nil { 115 | log.Fatal(err) 116 | response.Result = "error" 117 | response.Err = err 118 | } 119 | 120 | results <- response 121 | }(host) 122 | } 123 | } 124 | -------------------------------------------------------------------------------- /sessions/ssh.go: -------------------------------------------------------------------------------- 1 | package sessions 2 | 3 | import ( 4 | "bytes" 5 | "fmt" 6 | "log" 7 | 8 | "github.com/jmsdnns/omnitool/hosts" 9 | "golang.org/x/crypto/ssh" 10 | ) 11 | 12 | func dialServer(hostname string, config *ssh.ClientConfig) (*ssh.Client, error) { 13 | client, err := ssh.Dial("tcp", hostname, config) 14 | if err != nil { 15 | return nil, err 16 | } 17 | return client, nil 18 | } 19 | 20 | // SSHConnection is a container for the pieces necessary to hold an SSH connection 21 | // open in a goroutine 22 | type SSHConnection struct { 23 | hostname string 24 | config *ssh.ClientConfig 25 | client *ssh.Client 26 | session *ssh.Session 27 | } 28 | 29 | func (s *SSHConnection) init(hostname string, username string, keypath string) error { 30 | s.hostname = hostname 31 | 32 | // Instantiate config 33 | config, err := generateConfig(username, keypath) 34 | if err != nil { 35 | log.Fatal(err) 36 | return err 37 | } 38 | s.config = config 39 | 40 | // Establish connection 41 | client, err := dialServer(hostname, config) 42 | if err != nil { 43 | log.Fatal(err) 44 | return err 45 | } 46 | s.client = client 47 | 48 | // Make a session 49 | session, err := client.NewSession() 50 | if err != nil { 51 | log.Fatal(err) 52 | return err 53 | } 54 | s.session = session 55 | 56 | return nil 57 | } 58 | 59 | // TearDown closes all of the connections in an SSH connection 60 | func (s *SSHConnection) TearDown() { 61 | s.session.Close() 62 | s.client.Close() 63 | } 64 | 65 | func (s *SSHConnection) executeCmd(cmd string) string { 66 | var stdoutBuf bytes.Buffer 67 | s.session.Stdout = &stdoutBuf 68 | s.session.Run(cmd) 69 | 70 | return stdoutBuf.String() 71 | } 72 | 73 | // InitSSHConnection takes host details and returns an SSHConnection 74 | func InitSSHConnection(address string, username string, keypath string) (*SSHConnection, error) { 75 | sshConn := &SSHConnection{} 76 | err := sshConn.init(address, username, keypath) 77 | if err != nil { 78 | return nil, err 79 | } 80 | 81 | return sshConn, nil 82 | } 83 | 84 | // SSHResponse contains the result of running a command on a host via SSH 85 | type SSHResponse struct { 86 | Host string 87 | Result string 88 | Err error 89 | } 90 | 91 | // MapCmd takes the details for a command and maps it, via SSH, across a list 92 | // of hosts 93 | func MapCmd(hosts hosts.Group, username string, keypath string, command string, results chan SSHResponse) { 94 | for _, host := range hosts { 95 | go func(host string) { 96 | response := SSHResponse{Host: host} 97 | 98 | sshConn, err := InitSSHConnection(host, username, keypath) 99 | defer sshConn.TearDown() 100 | 101 | if err != nil { 102 | fmt.Println(err.Error()) 103 | response.Err = err 104 | } else { 105 | result := sshConn.executeCmd(command) 106 | response.Result = result 107 | } 108 | 109 | results <- response 110 | }(host) 111 | } 112 | } 113 | --------------------------------------------------------------------------------