From f5b8a5504a43f147a8991ec243cc40e172436d30 Mon Sep 17 00:00:00 2001 From: Karthik Dharmarajan Date: Mon, 31 Oct 2022 23:36:07 -0700 Subject: [PATCH] Addresses #97 via vendoring Former-commit-id: 09b2891dea2d06e57e1a1e9a0f042beec4a6ee6e --- fogros2/fogros2/vpn.py | 16 +- fogros2/fogros2/wgconfig.py | 287 ++++++++++++++++++++++++++++++++++++ fogros2/fogros2/wgexec.py | 90 +++++++++++ 3 files changed, 385 insertions(+), 8 deletions(-) create mode 100644 fogros2/fogros2/wgconfig.py create mode 100644 fogros2/fogros2/wgexec.py diff --git a/fogros2/fogros2/vpn.py b/fogros2/fogros2/vpn.py index bd0baac..4b7f1e8 100755 --- a/fogros2/fogros2/vpn.py +++ b/fogros2/fogros2/vpn.py @@ -33,8 +33,8 @@ import os -import wgconfig -import wgconfig.wgexec as wgexec +from .wgconfig import WGConfig +from .wgexec import get_publickey, generate_privatekey class VPN: @@ -49,8 +49,8 @@ def __init__( self.cloud_name_to_pub_key_path = dict() self.cloud_name_to_priv_key_path = dict() - self.robot_private_key = wgexec.generate_privatekey() - self.robot_public_key = wgexec.get_publickey(self.robot_private_key) + self.robot_private_key = generate_privatekey() + self.robot_public_key = get_publickey(self.robot_private_key) def generate_key_pairs(self, machines): """ @@ -60,9 +60,9 @@ def generate_key_pairs(self, machines): """ for machine in machines: name = machine.name - cloud_private_key = wgexec.generate_privatekey() + cloud_private_key = generate_privatekey() self.cloud_name_to_priv_key_path[name] = cloud_private_key - cloud_public_key = wgexec.get_publickey(cloud_private_key) + cloud_public_key = get_publickey(cloud_private_key) self.cloud_name_to_pub_key_path[name] = cloud_public_key def generate_wg_config_files(self, machines): @@ -74,7 +74,7 @@ def generate_wg_config_files(self, machines): name = machine.name machine_config_pwd = self.cloud_key_path + name machine_priv_key = self.cloud_name_to_priv_key_path[name] - aws_config = wgconfig.WGConfig(machine_config_pwd) + aws_config = WGConfig(machine_config_pwd) aws_config.add_attr(None, "PrivateKey", machine_priv_key) aws_config.add_attr(None, "ListenPort", 51820) aws_config.add_attr(None, "Address", f"10.0.0.{counter:d}/24") @@ -86,7 +86,7 @@ def generate_wg_config_files(self, machines): counter += 1 # generate robot configs - robot_config = wgconfig.WGConfig(self.robot_key_path) + robot_config = WGConfig(self.robot_key_path) robot_config.add_attr(None, "PrivateKey", self.robot_private_key) robot_config.add_attr(None, "ListenPort", 51820) robot_config.add_attr(None, "Address", "10.0.0.1/24") diff --git a/fogros2/fogros2/wgconfig.py b/fogros2/fogros2/wgconfig.py new file mode 100644 index 0000000..706c80f --- /dev/null +++ b/fogros2/fogros2/wgconfig.py @@ -0,0 +1,287 @@ +# Copyright 2022 The Regents of the University of California (Regents) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Copyright ©2022. The Regents of the University of California (Regents). +# All Rights Reserved. Permission to use, copy, modify, and distribute this +# software and its documentation for educational, research, and not-for-profit +# purposes, without fee and without a signed licensing agreement, is hereby +# granted, provided that the above copyright notice, this paragraph and the +# following two paragraphs appear in all copies, modifications, and +# distributions. Contact The Office of Technology Licensing, UC Berkeley, 2150 +# Shattuck Avenue, Suite 510, Berkeley, CA 94720-1620, (510) 643-7201, +# otl@berkeley.edu, http://ipira.berkeley.edu/industry-info for commercial +# licensing opportunities. IN NO EVENT SHALL REGENTS BE LIABLE TO ANY PARTY +# FOR DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, +# INCLUDING LOST PROFITS, ARISING OUT OF THE USE OF THIS SOFTWARE AND ITS +# DOCUMENTATION, EVEN IF REGENTS HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH +# DAMAGE. REGENTS SPECIFICALLY DISCLAIMS ANY WARRANTIES, INCLUDING, BUT NOT +# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A +# PARTICULAR PURPOSE. THE SOFTWARE AND ACCOMPANYING DOCUMENTATION, IF ANY, +# PROVIDED HEREUNDER IS PROVIDED "AS IS". REGENTS HAS NO OBLIGATION TO PROVIDE +# MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS. + + +from builtins import str +from builtins import range +from io import open +import os + + +class WGConfig(): + """A class for parsing and writing Wireguard configuration files""" + SECTION_FIRSTLINE = '_index_firstline' + SECTION_LASTLINE = '_index_lastline' + SECTION_RAW = '_rawdata' + _interface = None # interface attributes + _peers = None # peer data + + def __init__(self, file, keyattr='PublicKey'): + """Object initialization""" + self.filename = self.file2filename(file) + self.keyattr = keyattr + self.lines = [] + self.initialize_file() + + @staticmethod + def file2filename(file): + """Handle special filenames: 'wg0' and 'wg0.conf' become '/etc/wireguard/wg0.conf' """ + if os.path.basename(file) == file: + if not file.endswith('.conf'): + file += '.conf' + file = os.path.join('/etc/wireguard', file) + return file + + def invalidate_data(self): + """Clears the data structs""" + self._interface = None + self._peers = None + + def read_file(self): + """Reads the Wireguard config file into memory""" + with open(self.filename, 'r') as wgfile: + self.lines = [line.rstrip() for line in wgfile.readlines()] + self.invalidate_data() + + def write_file(self, file=None): + """Writes a Wireguard config file from memory to file""" + if file is None: + filename = self.filename + else: + filename = self.file2filename(file) + with os.fdopen(os.open(filename, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o640), 'w') as wgfile: + wgfile.writelines(line + '\n' for line in self.lines) + + @staticmethod + def parse_line(line): + """Splits a single attr/value line into its parts""" + attr, _, value = line.partition('=') + attr = attr.strip() + parts = value.partition('#') + value = parts[0].strip() # strip comments and whitespace + value = str(value) # this line is for Python2 support only + comment = parts[1] + parts[2] + if value.isnumeric(): + value = [int(value)] + else: + value = [item.strip() for item in value.split(',')] # decompose into list based on commata as separator + return attr, value, comment + + def parse_lines(self): + """Parses the lines of a Wireguard config file into memory""" + + # There will be two special attributes in the parsed data: + #_index_firstline: Line (zero indexed) of the section header (including any leading lines with comments) + #_index_lastline: Line (zero indexed) of the last attribute line of the section (including any directly following comments) + + def close_section(section, section_data): + section_data = {k: (v if len(v) > 1 else v[0]) for k, v in section_data.items()} + if section is None: # nothing to close on first section + return + elif section == 'interface': # close interface section + self._interface = section_data + else: # close peer section + peername = section_data.get(self.keyattr) + self._peers[peername] = section_data + section_data[self.SECTION_RAW] = self.lines[section_data[self.SECTION_FIRSTLINE]:(section_data[self.SECTION_LASTLINE] + 1)] + + self._interface = dict() + self._peers = dict() + section = None + section_data = dict() + last_empty_line_in_section = -1 # virtual empty line before start of file + for i, line in enumerate(self.lines): + # Ignore leading whitespace and trailing whitespace + line = line.strip() + # Ignore empty lines and comments + if len(line) == 0: + last_empty_line_in_section = i + continue + if line.startswith('['): # section + if last_empty_line_in_section is not None: + section_data[self.SECTION_LASTLINE] = [last_empty_line_in_section - 1] + close_section(section, section_data) + section_data = dict() + section = line[1:].partition(']')[0].lower() + if last_empty_line_in_section is None: + section_data[self.SECTION_FIRSTLINE] = [i] + else: + section_data[self.SECTION_FIRSTLINE] = [last_empty_line_in_section + 1] + last_empty_line_in_section = None + section_data[self.SECTION_LASTLINE] = [i] + if not section in ['interface', 'peer']: + raise ValueError('Unsupported section [{0}] in line {1}'.format(section, i)) + elif line.startswith('#'): + section_data[self.SECTION_LASTLINE] = [i] + else: # regular line + attr, value, _comment = self.parse_line(line) + section_data[attr] = section_data.get(attr, []) + section_data[attr].extend(value) + section_data[self.SECTION_LASTLINE] = [i] + close_section(section, section_data) + + def handle_leading_comment(self, leading_comment): + """Appends a leading comment for a section""" + if leading_comment is not None: + if leading_comment.strip()[0] != '#': + raise ValueError('A comment needs to start with a "#"') + self.lines.append(leading_comment) + + def initialize_file(self, leading_comment=None): + """Empties the file and adds the interface section header""" + self.lines = list() + self.handle_leading_comment(leading_comment) # add leading comment if needed + self.lines.append('[Interface]') + self.invalidate_data() + + def add_peer(self, key, leading_comment=None): + """Adds a new peer with the given (public) key""" + if key in self.peers: + raise KeyError('Peer to be added already exists') + self.lines.append('') # append an empty line for separation + self.handle_leading_comment(leading_comment) # add leading comment if needed + # Append peer with key attribute + self.lines.append('[Peer]') + self.lines.append('{0} = {1}'.format(self.keyattr, key)) + # Invalidate data cache + self.invalidate_data() + + def del_peer(self, key): + """Removes the peer with the given (public) key""" + if not key in self.peers: + raise KeyError('The peer to be deleted does not exist') + section_firstline = self.peers[key][self.SECTION_FIRSTLINE] + section_lastline = self.peers[key][self.SECTION_LASTLINE] + # Remove a blank line directly before the peer section + if section_firstline > 0: + if len(self.lines[section_firstline - 1]) == 0: + section_firstline -= 1 + # Only keep needed lines + result = [] + if section_firstline > 0: + result.extend(self.lines[0:section_firstline]) + result.extend(self.lines[(section_lastline + 1):]) + self.lines = result + # Invalidate data cache + self.invalidate_data() + + def get_sectioninfo(self, key): + """Get first and last line of the section identified by the given key ("None" for interface section)""" + if key is None: # interface + section_firstline = self.interface[self.SECTION_FIRSTLINE] + section_lastline = self.interface[self.SECTION_LASTLINE] + else: # peer + if not key in self.peers: + raise KeyError('The specified peer does not exist') + section_firstline = self.peers[key][self.SECTION_FIRSTLINE] + section_lastline = self.peers[key][self.SECTION_LASTLINE] + return section_firstline, section_lastline + + def add_attr(self, key, attr, value, leading_comment=None, append_as_line=False): + """Adds an attribute/value pair to the given peer ("None" for adding an interface attribute)""" + section_firstline, section_lastline = self.get_sectioninfo(key) + if leading_comment is not None: + if leading_comment.strip()[0] != '#': + raise ValueError('A comment needs to start with a "#"') + # Look for line with the attribute + line_found = None + for i in range(section_firstline + 1, section_lastline + 1): + line_attr, line_value, line_comment = self.parse_line(self.lines[i]) + if attr == line_attr: + line_found = i + # Add the attribute at the right place + if (line_found is None) or append_as_line: + line_found = section_lastline if (line_found is None) else line_found + line_found += 1 + self.lines.insert(line_found, '{0} = {1}'.format(attr, value)) + else: + line_attr, line_value, line_comment = self.parse_line(self.lines[line_found]) + line_value.append(value) + if len(line_comment) > 0: + line_comment = ' ' + line_comment + line_value = [str(item) for item in line_value] + self.lines[line_found] = line_attr + ' = ' + ', '.join(line_value) + line_comment + # Handle leading comments + if leading_comment is not None: + self.lines.insert(line_found, leading_comment) + # Invalidate data cache + self.invalidate_data() + + def del_attr(self, key, attr, value=None, remove_leading_comments=True): + """Removes an attribute/value pair from the given peer ("None" for adding an interface attribute); set 'value' to 'None' to remove all values""" + section_firstline, section_lastline = self.get_sectioninfo(key) + # Find all lines with matching attribute name and (if requested) value + line_found = [] + for i in range(section_firstline + 1, section_lastline + 1): + line_attr, line_value, line_comment = self.parse_line(self.lines[i]) + if attr == line_attr: + if (value is None) or (value in line_value): + line_found.append(i) + if len(line_found) == 0: + raise ValueError('The attribute/value to be deleted is not present') + # Process all relevant lines + for i in reversed(line_found): # reversed so that non-processed indices stay valid + if value is None: + del(self.lines[i]) + else: + line_attr, line_value, line_comment = self.parse_line(self.lines[i]) + line_value.remove(value) + if len(line_value) > 0: # keep remaining values in that line + self.lines[i] = line_attr + ' = ' + ', '.join(line_value) + line_comment + else: # otherwise line is no longer needed + del(self.lines[i]) + # Handle leading comments + if remove_leading_comments: + i = line_found[0] - 1 + while i > 0: + if len(self.lines[i]) and (self.lines[i][0] == '#'): + del(self.lines[i]) + i -= 1 + else: + break + # Invalidate data cache + self.invalidate_data() + + @property + def interface(self): + """Dictionary with interface attributes""" + if self._interface is None: + self.parse_lines() + return self._interface + + @property + def peers(self): + """Dictionary with peer data""" + if self._peers is None: + self.parse_lines() + return self._peers \ No newline at end of file diff --git a/fogros2/fogros2/wgexec.py b/fogros2/fogros2/wgexec.py new file mode 100644 index 0000000..fffe250 --- /dev/null +++ b/fogros2/fogros2/wgexec.py @@ -0,0 +1,90 @@ +# Copyright 2022 The Regents of the University of California (Regents) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Copyright ©2022. The Regents of the University of California (Regents). +# All Rights Reserved. Permission to use, copy, modify, and distribute this +# software and its documentation for educational, research, and not-for-profit +# purposes, without fee and without a signed licensing agreement, is hereby +# granted, provided that the above copyright notice, this paragraph and the +# following two paragraphs appear in all copies, modifications, and +# distributions. Contact The Office of Technology Licensing, UC Berkeley, 2150 +# Shattuck Avenue, Suite 510, Berkeley, CA 94720-1620, (510) 643-7201, +# otl@berkeley.edu, http://ipira.berkeley.edu/industry-info for commercial +# licensing opportunities. IN NO EVENT SHALL REGENTS BE LIABLE TO ANY PARTY +# FOR DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, +# INCLUDING LOST PROFITS, ARISING OUT OF THE USE OF THIS SOFTWARE AND ITS +# DOCUMENTATION, EVEN IF REGENTS HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH +# DAMAGE. REGENTS SPECIFICALLY DISCLAIMS ANY WARRANTIES, INCLUDING, BUT NOT +# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A +# PARTICULAR PURPOSE. THE SOFTWARE AND ACCOMPANYING DOCUMENTATION, IF ANY, +# PROVIDED HEREUNDER IS PROVIDED "AS IS". REGENTS HAS NO OBLIGATION TO PROVIDE +# MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS. + + +import logging +import shlex +import subprocess + + +logger = logging.getLogger(__name__) + + +def execute(command, input=None, suppressoutput=False, suppresserrors=False): + """Execute a command""" + args = shlex.split(command) + stdin = None if input is None else subprocess.PIPE + input = None if input is None else input.encode('utf-8') + nsp = subprocess.Popen(args, stdin=stdin, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + out, err = nsp.communicate(input=input) + if err is not None: + err = err.decode('utf8') + if not suppresserrors and (len(err) > 0): + logger.error(err) + out = out.decode('utf8') + if not suppressoutput and (len(out) > 0): + print(out) + nsp.wait() + return out, err, nsp.returncode + +def generate_privatekey(): + """Generates a WireGuard private key""" + out, err, returncode = execute('wg genkey', suppressoutput=True) + if (returncode != 0) or (len(err) > 0): + return None + out = out.strip() # remove trailing newline + return out + +def get_publickey(wg_private): + """Gets the public key belonging to the given WireGuard private key""" + if wg_private is None: + return None + out, err, returncode = execute('wg pubkey', input=wg_private, suppressoutput=True) + if (returncode != 0) or (len(err) > 0): + return None + out = out.strip() # remove trailing newline + return out + +def generate_keypair(): + """Generates a WireGuard key pair (returns tuple of private key and public key)""" + wg_private = generate_privatekey() + wg_public = get_publickey(wg_private) + return wg_private, wg_public + +def generate_presharedkey(): + """Generates a WireGuard preshared key""" + out, err, returncode = execute('wg genpsk', suppressoutput=True) + if (returncode != 0) or (len(err) > 0): + return None + out = out.strip() # remove trailing newline + return out \ No newline at end of file