<commit_before>#!/usr/bin/python

# Copyright (c) 2016 Hector Martin <marcan@marcan.st>
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

from __future__ import print_function

import sys, os.path, struct, argparse
from xmodem import XMODEM
from serial import Serial

BIN_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "bin")

class MtkFlasher(object):
    BAUD = {
        "ls": 115200,
        "hs": 921600,
        "ss": 3000000,
    }
    MAGIC = 0x00052001

    def __init__(self, port, speed, binpath=None, debug=False):
        self.speed = speed
        self.baud = self.BAUD[speed]
        self.port = port
        self.xm = XMODEM(self.getc, self.putc, "xmodem1k", "\xff")
        self.binpath = binpath or BIN_PATH
        self.tag = 0x1000
        self.debug = debug

    def log(self, *args):
        if self.debug:
            print(*args)

    def getc(self, size, timeout=1):
        self.ser.timeout = timeout
        return self.ser.read(size)

    def putc(self, data, timeout=1):
        self.ser.write_timeout = timeout
        return self.ser.write(data)

    def xm_send(self, fd, msg="  Sending"):
        fd.seek(0, 2)
        length = fd.tell()
        fd.seek(0, 0)
        blocks = (length + 1023) / 1024

        def callback(total, success, error):
            sys.stdout.write("\r%s... %d%%" % (msg, 100 * success // blocks))
            sys.stdout.flush()

        self.xm.send(fd, callback=callback)
        print()

    def bootstrap(self, ):
        print("Opening port at 115200 baud...")
        self.ser = Serial(self.port, 115200)
        fname = "uart_%s.bin" % self.speed
        print("Sending baudrate switcher (%s)..." % fname)
        self.xm_send(open(os.path.join(self.binpath, fname), "rb"))
        self.ser.close()

        print("Reopening port at %d baud..." % self.baud)
        self.ser = Serial(self.port, self.baud)
        fname = "ated_%s.bin" % self.speed
        print("Sending ATED (%s)..." % fname)
        self.xm_send(open(os.path.join(self.binpath, fname), "rb"))

    def command(self, cmd, args=b""):
        self.tag += 1
        self.ser.timeout = 5
        self.ser.write_timeout = 5

        payload = struct.pack(">H", cmd) + args
        data = struct.pack(">IHH", 0x00052001,
                           len(payload) + 2, self.tag) + payload
        crc = self.xm.calc_crc(data)
        self.ser.write(data + struct.pack(">H", crc))

        reply = self.ser.read(8)
        magic, reply_len, tag = struct.unpack(">IHH", reply)
        assert magic == (self.MAGIC | 0x80000000)
        assert tag == self.tag
        reply += self.ser.read(reply_len)
        assert len(reply) == reply_len + 8
        crc = struct.unpack(">H", reply[-2:])[0]
        assert self.xm.calc_crc(reply[:-2]) == crc
        reply_cmd = struct.unpack(">H", reply[8:10])[0]
        assert reply_cmd == (cmd + 1)
        return reply[10:-2]

    def initialize(self):
        self.log("  cmd: initialize")
        reply = self.command(0x00)
        self.log("    returned:", reply.encode("hex"))

    def get_storage_info(self):
        self.log("  cmd: get_storage_info")
        reply = self.command(0x10)
        self.log("    returned:", reply.encode("hex"))

        unk1, unk2, flash_size, = struct.unpack(">III", reply)
        return {"size": flash_size}

    def erase(self, start, length):
        self.log("  cmd: erase(0x%x, 0x%x)" % (start, length))
        reply = self.command(0x0a, struct.pack(">II", start, length))
        self.log("    returned:", reply.encode("hex"))

    def erase_end(self):
        self.log("  cmd: erase_end")
        reply = self.command(0x0c)
        self.log("    returned:", reply.encode("hex"))

    def download(self, address, length):
        self.log("  cmd: download(0x%x, 0x%x)" % (address, length))
        flags = 0x01000400
        reply = self.command(0x02, struct.pack(">IIII",
                                               address, length,
                                               address + length, flags))
        self.log("    returned:", reply.encode("hex"))

    def download_end(self):
        self.log("  cmd: download_end")
        reply = self.command(0x04)
        self.log("    returned:", reply.encode("hex"))

    # Orig uses 0x20000 but this provides more granular progress
    def erase_range(self, addr, length, blocksize=0x2000):
        for i in xrange(0, length, blocksize):
            self.erase(addr + i, min(blocksize, length - i))
            if not self.debug:
                sys.stdout.write("\r  Erasing... %d%%" % (100 * i // length))
                sys.stdout.flush()
        self.erase_end()
        if not self.debug:
            print("\r  Erasing... 100%")

    def write_file(self, addr, filename, erase=True):
        print("Writing to 0x%x: %s" % (addr, filename))
        fd = open(filename, "rb")
        fd.seek(0, 2)
        length = fd.tell()
        fd.seek(0, 0)

        assert addr & 0xfff == 0
        padded_length = (length + 0xfff) & ~0xfff

        if erase:
            self.erase_range(addr, padded_length)

        self.download(addr, length)
        self.xm_send(fd, "  Writing")
        self.download_end()

if __name__ == "__main__":
    def parse_write(s):
        off, f = s.split(":", 1)
        off = int(off, 0)
        return off, f

    parser = argparse.ArgumentParser()
    parser.add_argument("-v", "--verbose", action="store_true",
                        help="show debug output")
    parser.add_argument("-b", "--bin",
                        help="path to the directory containing flasher blobs")
    parser.add_argument("-p", "--port", metavar="PORT", default="/dev/ttyACM0",
                        help="serial port device")
    parser.add_argument("-s", "--speed", metavar="SPEED", default="hs",
                        help="port speed (ls|hs|ss)", choices=["ls","hs","ss"])
    parser.add_argument("-e", "--erase", action="store_true",
                        help="wipe the entire flash memory first")
    parser.add_argument("-w", "--write", metavar="ADDR:FILE", action="append",
                        help="write a file to flash memory", type=parse_write)

    args = parser.parse_args()

    mtk = MtkFlasher(args.port, args.speed, binpath=args.bin,
                     debug=args.verbose)

    mtk.bootstrap()
    mtk.initialize()
    info = mtk.get_storage_info()
    print("Flash size: 0x%x" % info["size"])

    if args.erase:
        print("Erasing flash memory...")
        mtk.erase_range(0, info["size"])

    if args.write:
        for off, filename in args.write:
            mtk.write_file(off, filename, erase=not args.erase)

<commit_msg>Add command to read eFuses<commit_after>#!/usr/bin/python

# Copyright (c) 2016 Hector Martin <marcan@marcan.st>
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

from __future__ import print_function

import sys, os.path, struct, argparse
from xmodem import XMODEM
from serial import Serial

BIN_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "bin")

class MtkFlasher(object):
    BAUD = {
        "ls": 115200,
        "hs": 921600,
        "ss": 3000000,
    }
    MAGIC = 0x00052001
    SECRET_MAGIC = 0x00052002

    def __init__(self, port, speed, binpath=None, debug=False):
        self.speed = speed
        self.baud = self.BAUD[speed]
        self.port = port
        self.xm = XMODEM(self.getc, self.putc, "xmodem1k", "\xff")
        self.binpath = binpath or BIN_PATH
        self.tag = 0x1000
        self.debug = debug

    def log(self, *args):
        if self.debug:
            print(*args)

    def getc(self, size, timeout=1):
        self.ser.timeout = timeout
        return self.ser.read(size)

    def putc(self, data, timeout=1):
        self.ser.write_timeout = timeout
        return self.ser.write(data)

    def xm_send(self, fd, msg="  Sending"):
        fd.seek(0, 2)
        length = fd.tell()
        fd.seek(0, 0)
        blocks = (length + 1023) / 1024

        def callback(total, success, error):
            sys.stdout.write("\r%s... %d%%" % (msg, 100 * success // blocks))
            sys.stdout.flush()

        self.xm.send(fd, callback=callback)
        print()

    def bootstrap(self, ):
        print("Opening port at 115200 baud...")
        self.ser = Serial(self.port, 115200)
        fname = "uart_%s.bin" % self.speed
        print("Sending baudrate switcher (%s)..." % fname)
        self.xm_send(open(os.path.join(self.binpath, fname), "rb"))
        self.ser.close()

        print("Reopening port at %d baud..." % self.baud)
        self.ser = Serial(self.port, self.baud)
        fname = "ated_%s.bin" % self.speed
        print("Sending ATED (%s)..." % fname)
        self.xm_send(open(os.path.join(self.binpath, fname), "rb"))

    def command(self, cmd, args=b"", magic=MAGIC):
        self.tag += 1
        self.ser.timeout = 5
        self.ser.write_timeout = 5

        payload = struct.pack(">H", cmd) + args
        data = struct.pack(">IHH", magic,
                           len(payload) + 2, self.tag) + payload
        crc = self.xm.calc_crc(data)
        self.ser.write(data + struct.pack(">H", crc))

        reply = self.ser.read(8)
        reply_magic, reply_len, tag = struct.unpack(">IHH", reply)
        assert reply_magic == (magic | 0x80000000)
        assert tag == self.tag
        reply += self.ser.read(reply_len)
        assert len(reply) == reply_len + 8
        crc = struct.unpack(">H", reply[-2:])[0]
        assert self.xm.calc_crc(reply[:-2]) == crc
        reply_cmd = struct.unpack(">H", reply[8:10])[0]
        assert reply_cmd == (cmd + 1)
        return reply[10:-2]

    def initialize(self):
        self.log("  cmd: initialize")
        reply = self.command(0x00)
        self.log("    returned:", reply.encode("hex"))

    def get_storage_info(self):
        self.log("  cmd: get_storage_info")
        reply = self.command(0x10)
        self.log("    returned:", reply.encode("hex"))

        unk1, unk2, flash_size, = struct.unpack(">III", reply)
        return {"size": flash_size}

    def erase(self, start, length):
        self.log("  cmd: erase(0x%x, 0x%x)" % (start, length))
        reply = self.command(0x0a, struct.pack(">II", start, length))
        self.log("    returned:", reply.encode("hex"))

    def erase_end(self):
        self.log("  cmd: erase_end")
        reply = self.command(0x0c)
        self.log("    returned:", reply.encode("hex"))

    def download(self, address, length):
        self.log("  cmd: download(0x%x, 0x%x)" % (address, length))
        flags = 0x01000400
        reply = self.command(0x02, struct.pack(">IIII",
                                               address, length,
                                               address + length, flags))
        self.log("    returned:", reply.encode("hex"))

    def download_end(self):
        self.log("  cmd: download_end")
        reply = self.command(0x04)
        self.log("    returned:", reply.encode("hex"))

    # Orig uses 0x20000 but this provides more granular progress
    def erase_range(self, addr, length, blocksize=0x2000):
        for i in xrange(0, length, blocksize):
            self.erase(addr + i, min(blocksize, length - i))
            if not self.debug:
                sys.stdout.write("\r  Erasing... %d%%" % (100 * i // length))
                sys.stdout.flush()
        self.erase_end()
        if not self.debug:
            print("\r  Erasing... 100%")

    def write_file(self, addr, filename, erase=True):
        print("Writing to 0x%x: %s" % (addr, filename))
        fd = open(filename, "rb")
        fd.seek(0, 2)
        length = fd.tell()
        fd.seek(0, 0)

        assert addr & 0xfff == 0
        padded_length = (length + 0xfff) & ~0xfff

        if erase:
            self.erase_range(addr, padded_length)

        self.download(addr, length)
        self.xm_send(fd, "  Writing")
        self.download_end()

    def read_efuse(self, off, length):
        assert (off & 0xf == 0)
        assert (length & 0xf == 0)
        self.log("  cmd: read_efuse")
        reply = self.command(0x04, struct.pack(">II", off, length),
                             self.SECRET_MAGIC)
        self.log("    returned:", reply.encode("hex"))
        return reply[4:]

    def get_mac(self):
        fuseblock = self.read_efuse(0, 0x10)
        return ":".join("%02x" % ord(c) for c in fuseblock[4:10])

if __name__ == "__main__":
    def parse_write(s):
        off, f = s.split(":", 1)
        off = int(off, 0)
        return off, f

    parser = argparse.ArgumentParser()
    parser.add_argument("-v", "--verbose", action="store_true",
                        help="show debug output")
    parser.add_argument("-b", "--bin",
                        help="path to the directory containing flasher blobs")
    parser.add_argument("-p", "--port", metavar="PORT", default="/dev/ttyACM0",
                        help="serial port device")
    parser.add_argument("-s", "--speed", metavar="SPEED", default="hs",
                        help="port speed (ls|hs|ss)", choices=["ls","hs","ss"])
    parser.add_argument("-e", "--erase", action="store_true",
                        help="wipe the entire flash memory first")
    parser.add_argument("-w", "--write", metavar="ADDR:FILE", action="append",
                        help="write a file to flash memory", type=parse_write)
    parser.add_argument("-f", "--read-fuses", metavar="FILE",
                        help="read eFuses to a file")

    args = parser.parse_args()

    mtk = MtkFlasher(args.port, args.speed, binpath=args.bin,
                     debug=args.verbose)

    mtk.bootstrap()
    mtk.initialize()
    info = mtk.get_storage_info()
    print("Flash size: 0x%x" % info["size"])
    print("MAC address: %s" % mtk.get_mac())

    if args.read_fuses:
        fuses = mtk.read_efuse(0, 0x200)
        with open(args.read_fuses, "wb") as fd:
            fd.write(fuses)

    if args.erase:
        print("Erasing flash memory...")
        mtk.erase_range(0, info["size"])

    if args.write:
        for off, filename in args.write:
            mtk.write_file(off, filename, erase=not args.erase)

<|endoftext|>


<commit_before>from math import cos, sin, pi
from numba import njit

@njit
def hard_cutoff(r_cut, ri, ci):
    return 1, 0

@njit
def quadratic_cutoff(r_cut, ri, ci):
    rdiff = r_cut - ri
    fi = rdiff * rdiff
    fdi = 2 * rdiff * ci

    return fi, fdi

@njit
def cubic_cutoff(r_cut, ri, ci):
    rdiff = r_cut - ri
    fi = rdiff * rdiff * rdiff
    fdi = 3 * rdiff * rdiff * ci

    return fi, fdi

@njit
def cosine_cutoff(r_cut, ri, ci, d=1):
    if ri > r_cut - d:
        fi = (1/2) * (cos(pi * (ri - r_cut + d) / d) + 1)
        fdi = (pi/(2 * d)) * sin(pi * (r_cut - ri) / d) * ci
    else:
        fi = 1
        fdi = 0

    return fi, fdi
<commit_msg>add cutoff docstrings<commit_after>"""
The cutoffs module gives a few different options for smoothly sending the GP
kernel to 0 near the boundary of the cutoff sphere.
"""
from math import cos, sin, pi
from numba import njit

@njit
def hard_cutoff(r_cut, ri, ci):
    """A hard cutoff that assigns a value of 1 to all interatomic distances.

    :param r_cut: Cutoff value (in angstrom).
    :type r_cut: float
    :param ri: Interatomic distance.
    :type ri: float
    :param ci: Cartesian coordinate divided by the distance.
    :type ci: float
    :return: Cutoff value and its derivative.
    :rtype: float, float
    """
    return 1, 0

@njit
def quadratic_cutoff(r_cut, ri, ci):
    """A quadratic cutoff that goes to zero smoothly at the cutoff boundary.

    :param r_cut: Cutoff value (in angstrom).
    :type r_cut: float
    :param ri: Interatomic distance.
    :type ri: float
    :param ci: Cartesian coordinate divided by the distance.
    :type ci: float
    :return: Cutoff value and its derivative.
    :rtype: float, float
    """
    rdiff = r_cut - ri
    fi = rdiff * rdiff
    fdi = 2 * rdiff * ci

    return fi, fdi

@njit
def cubic_cutoff(r_cut, ri, ci):
    """A cubic cutoff that goes to zero smoothly at the cutoff boundary.

    :param r_cut: Cutoff value (in angstrom).
    :type r_cut: float
    :param ri: Interatomic distance.
    :type ri: float
    :param ci: Cartesian coordinate divided by the distance.
    :type ci: float
    :return: Cutoff value and its derivative.
    :rtype: float, float
    """

    rdiff = r_cut - ri
    fi = rdiff * rdiff * rdiff
    fdi = 3 * rdiff * rdiff * ci

    return fi, fdi

@njit
def cosine_cutoff(r_cut, ri, ci, d=1):
    """A cosine cutoff that returns 1 up to r_cut - d, and assigns a cosine envelope to values of r between (r_cut - d) and r_cut.

    :param r_cut: Cutoff value (in angstrom).
    :type r_cut: float
    :param ri: Interatomic distance.
    :type ri: float
    :param ci: Cartesian coordinate divided by the distance.
    :type ci: float
    :return: Cutoff value and its derivative.
    :rtype: float, float
    """

    if ri > r_cut - d:
        fi = (1/2) * (cos(pi * (ri - r_cut + d) / d) + 1)
        fdi = (pi/(2 * d)) * sin(pi * (r_cut - ri) / d) * ci
    else:
        fi = 1
        fdi = 0

    return fi, fdi
<|endoftext|>


<commit_before>#!/usr/bin/env python
from __future__ import print_function

import re
import sys
from copy import copy
from time import time
from uuid import uuid4

__version__ = "1.1.1"

methods_order = ["GET", "POST", "PUT", "PATCH", "DELETE", "COPY", "HEAD",
                 "OPTIONS", "LINK", "UNLINK", "PURGE"]

var_re = re.compile(r"(?P<var><[a-zA-Z0-9_]+:(?P<var_name>[a-zA-Z0-9_]+)>)")

venv_warning = ("WARNING: Attempting to work in a virtualenv. If you encounter "
                "problems, please install flask2postman inside the virtualenv.")

PY2 = sys.version_info[0] < 3
if PY2:
    maxint = sys.maxint
else:
    maxint = sys.maxsize

def get_time():
    return int(round(time() * 1000))

class Collection:

    def __init__(self, name):
        self.id = str(uuid4())
        self._requests = []
        self._folders = []
        self.name = name
        self.timestamp = get_time()

    def reorder_requests(self):
        def _get_key(request):
            return str(methods_order.index(request.method)) + request.name
        self._requests = sorted(self._requests, key=_get_key)

    def add_folder(self, folder):
        folder._collection = self
        folder.collection_id = self.id
        folder.collection_name = self.name
        self._folders.append(folder)
        self.reorder_requests()

    def add_route(self, route):
        route.collection_id = self.id
        self._requests.append(route)
        self.reorder_requests()

    @property
    def order(self):
        return [request.id for request in self._requests if request.folder is None]

    @property
    def requests(self):
        return [request.to_dict() for request in self._requests]

    @property
    def folders(self):
        return [folder.to_dict() for folder in self._folders]

    def to_dict(self):
        d = copy(self.__dict__)
        d.pop("_requests")
        d.pop("_folders")
        d.update(requests=self.requests, order=self.order)
        if len(self.folders) > 0:
            d.update(folders=self.folders)
        return d

class Folder:
    _folders_ = dict()

    def __init__(self, name):
        self.id = str(uuid4())
        self._collection = None
        self._requests = []
        self.name = name

    def add_route(self, route):
        route.folder = self.id
        self._collection.add_route(route)
        self._requests.append(route)
        self.reorder_requests()

    def reorder_requests(self):
        def _get_key(request):
            return str(methods_order.index(request.method)) + request.name
        self._requests = sorted(self._requests, key=_get_key)

    @property
    def order(self):
        return [request.id for request in self._requests]

    def to_dict(self):
        d = copy(self.__dict__)
        d.pop("_requests")
        d.pop("_collection")
        d["collectionId"] = d.pop("collection_id")
        d["collectionName"] = d.pop("collection_name")
        d.update(order=self.order)
        return d

    @staticmethod
    def get_or_create_folder(collection, name):
        folder = Folder._folders_.get(name)
        if folder is not None:
            return folder
        else:
            folder = Folder(name)
            collection.add_folder(folder)
            Folder._folders_[name] = folder
            return folder

class Route:

    def __init__(self, name, url, method, description="", headers="", data=[],
                 data_mode="params"):
        self.id = str(uuid4())
        self.data = data
        self.data_mode = data_mode
        self.description = description
        self.headers = headers
        self.method = method
        self.name = name
        self.folder = None
        self.time = get_time()
        self.url = url

    def to_dict(self):
        d = copy(self.__dict__)
        d["collectionId"] = d.pop("collection_id")
        d["dataMode"] = d.pop("data_mode")
        if self.folder is None:
            d.pop("folder")
        return d

    @classmethod
    def from_werkzeug(cls, rule, method, base_url):
        name = rule.endpoint.rsplit('.', 1)[-1]
        name = name.split("_", 1)[-1]
        name = name.replace("_", " ")

        url = base_url + rule.rule
        for match in re.finditer(var_re, url):
            var = match.group("var")
            var_name = "{{" + match.group("var_name") + "}}"
            url = url.replace(var, var_name)
        return cls(name, url, method)

# ramnes: shamelessly stolen from https://www.python.org/dev/peps/pep-0257/
def trim(docstring):
    if not docstring:
        return ""
    lines = docstring.expandtabs().splitlines()
    indent = maxint
    for line in lines[1:]:
        stripped = line.lstrip()
        if stripped:
            indent = min(indent, len(line) - len(stripped))
    trimmed = [lines[0].strip()]
    if indent < maxint:
        for line in lines[1:]:
            trimmed.append(line[indent:].rstrip())
    while trimmed and not trimmed[-1]:
        trimmed.pop()
    while trimmed and not trimmed[0]:
        trimmed.pop(0)
    return '\n'.join(trimmed)

def main():
    import json
    import logging
    import os
    import site
    from argparse import ArgumentParser

    from flask import Flask, current_app

    sys.path.append(os.getcwd())

    venv = os.environ.get("VIRTUAL_ENV", None)
    if venv:
        print(venv_warning, file=sys.stderr)
        if sys.platform == "win32":
            path = os.path.join(venv, 'Lib', 'site-packages')
        else:
            python = "python{}.{}".format(*sys.version_info[:2])
            path = os.path.join(venv, 'lib', python, 'site-packages')
        sys.path.insert(0, path)
        site.addsitedir(path)

    parser = ArgumentParser()
    parser.add_argument("flask_instance")
    parser.add_argument("-n", "--name", default=os.path.basename(os.getcwd()),
                        help="Postman collection name (default: current directory name)")
    parser.add_argument("-b", "--base_url", default="{{base_url}}",
                        help="the base of every URL (default: {{base_url}})")
    parser.add_argument("-a", "--all", action="store_true",
                        help="also generate OPTIONS/HEAD methods")
    parser.add_argument("-i", "--indent", action="store_true",
                        help="indent the output")
    parser.add_argument("-f", "--folders", action="store_true",
                        help="add Postman folders for blueprints")
    args = parser.parse_args()

    logging.disable(logging.CRITICAL)

    try:
        app_path, app_name = args.flask_instance.rsplit('.', 1)
        app = getattr(__import__(app_path), app_name)
    except Exception as e:
        parser.error("can't import \"{}\" ({})".format(args.flask_instance, str(e)))

    if not isinstance(app, Flask):
        parser.error("\"{}\" is not a Flask instance".format(args.flask_instance))

    with app.app_context():
        collection = Collection(args.name)
        rules = list(current_app.url_map.iter_rules())
        for rule in rules:
            for method in rule.methods:
                if args.all or method not in ["OPTIONS", "HEAD"]:
                    endpoint = current_app.view_functions[rule.endpoint]
                    route = Route.from_werkzeug(rule, method, args.base_url)
                    route.description = trim(endpoint.__doc__)
                    folder_name = rule.endpoint.rsplit('.', maxsplit=2)
                    if len(folder_name) >= 2 and args.folders:
                        folder = Folder.get_or_create_folder(collection, folder_name[0])
                        folder.add_route(route)
                    else:
                        collection.add_route(route)

    if args.indent:
        json = json.dumps(collection.to_dict(), indent=4, sort_keys=True)
    else:
        json = json.dumps(collection.to_dict())

    print(json)

if __name__ == "__main__":
    main()
<commit_msg>Remove `copy.copy` in favor of a dict comprehension<commit_after>#!/usr/bin/env python
from __future__ import print_function

import re
import sys
from time import time
from uuid import uuid4

__version__ = "1.1.1"

methods_order = ["GET", "POST", "PUT", "PATCH", "DELETE", "COPY", "HEAD",
                 "OPTIONS", "LINK", "UNLINK", "PURGE"]

var_re = re.compile(r"(?P<var><[a-zA-Z0-9_]+:(?P<var_name>[a-zA-Z0-9_]+)>)")

venv_warning = ("WARNING: Attempting to work in a virtualenv. If you encounter "
                "problems, please install flask2postman inside the virtualenv.")

PY2 = sys.version_info[0] < 3
if PY2:
    maxint = sys.maxint
else:
    maxint = sys.maxsize

def get_time():
    return int(round(time() * 1000))

class Collection:

    def __init__(self, name):
        self.id = str(uuid4())
        self._requests = []
        self._folders = []
        self.name = name
        self.timestamp = get_time()

    def reorder_requests(self):
        def _get_key(request):
            return str(methods_order.index(request.method)) + request.name
        self._requests = sorted(self._requests, key=_get_key)

    def add_folder(self, folder):
        folder._collection = self
        folder.collection_id = self.id
        folder.collection_name = self.name
        self._folders.append(folder)
        self.reorder_requests()

    def add_route(self, route):
        route.collection_id = self.id
        self._requests.append(route)
        self.reorder_requests()

    @property
    def order(self):
        return [request.id for request in self._requests if request.folder is None]

    @property
    def requests(self):
        return [request.to_dict() for request in self._requests]

    @property
    def folders(self):
        return [folder.to_dict() for folder in self._folders]

    def to_dict(self):
        d = {k: v for k, v in self.__dict__.items() if not k.startswith("_")}
        d.update(requests=self.requests, order=self.order, folders=self.folders)
        return d

class Folder:
    _folders_ = dict()

    def __init__(self, name):
        self.id = str(uuid4())
        self._collection = None
        self._requests = []
        self.name = name

    def add_route(self, route):
        route.folder = self.id
        self._collection.add_route(route)
        self._requests.append(route)
        self.reorder_requests()

    def reorder_requests(self):
        def _get_key(request):
            return str(methods_order.index(request.method)) + request.name
        self._requests = sorted(self._requests, key=_get_key)

    @property
    def order(self):
        return [request.id for request in self._requests]

    def to_dict(self):
        d = {k: v for k, v in self.__dict__.items() if not k.startswith("_")}
        d["collectionId"] = d.pop("collection_id")
        d["collectionName"] = d.pop("collection_name")
        d.update(order=self.order)
        return d

    @staticmethod
    def get_or_create_folder(collection, name):
        folder = Folder._folders_.get(name)
        if folder is not None:
            return folder
        else:
            folder = Folder(name)
            collection.add_folder(folder)
            Folder._folders_[name] = folder
            return folder

class Route:

    def __init__(self, name, url, method, description="", headers="", data=[],
                 data_mode="params"):
        self.id = str(uuid4())
        self.data = data
        self.data_mode = data_mode
        self.description = description
        self.headers = headers
        self.method = method
        self.name = name
        self.folder = None
        self.time = get_time()
        self.url = url

    def to_dict(self):
        d = {k: v for k, v in self.__dict__.items() if not k.startswith("_")}
        d["collectionId"] = d.pop("collection_id")
        d["dataMode"] = d.pop("data_mode")
        if self.folder is None:
            d.pop("folder")
        return d

    @classmethod
    def from_werkzeug(cls, rule, method, base_url):
        name = rule.endpoint.rsplit('.', 1)[-1]
        name = name.split("_", 1)[-1]
        name = name.replace("_", " ")

        url = base_url + rule.rule
        for match in re.finditer(var_re, url):
            var = match.group("var")
            var_name = "{{" + match.group("var_name") + "}}"
            url = url.replace(var, var_name)
        return cls(name, url, method)

# ramnes: shamelessly stolen from https://www.python.org/dev/peps/pep-0257/
def trim(docstring):
    if not docstring:
        return ""
    lines = docstring.expandtabs().splitlines()
    indent = maxint
    for line in lines[1:]:
        stripped = line.lstrip()
        if stripped:
            indent = min(indent, len(line) - len(stripped))
    trimmed = [lines[0].strip()]
    if indent < maxint:
        for line in lines[1:]:
            trimmed.append(line[indent:].rstrip())
    while trimmed and not trimmed[-1]:
        trimmed.pop()
    while trimmed and not trimmed[0]:
        trimmed.pop(0)
    return '\n'.join(trimmed)

def main():
    import json
    import logging
    import os
    import site
    from argparse import ArgumentParser

    from flask import Flask, current_app

    sys.path.append(os.getcwd())

    venv = os.environ.get("VIRTUAL_ENV", None)
    if venv:
        print(venv_warning, file=sys.stderr)
        if sys.platform == "win32":
            path = os.path.join(venv, 'Lib', 'site-packages')
        else:
            python = "python{}.{}".format(*sys.version_info[:2])
            path = os.path.join(venv, 'lib', python, 'site-packages')
        sys.path.insert(0, path)
        site.addsitedir(path)

    parser = ArgumentParser()
    parser.add_argument("flask_instance")
    parser.add_argument("-n", "--name", default=os.path.basename(os.getcwd()),
                        help="Postman collection name (default: current directory name)")
    parser.add_argument("-b", "--base_url", default="{{base_url}}",
                        help="the base of every URL (default: {{base_url}})")
    parser.add_argument("-a", "--all", action="store_true",
                        help="also generate OPTIONS/HEAD methods")
    parser.add_argument("-i", "--indent", action="store_true",
                        help="indent the output")
    parser.add_argument("-f", "--folders", action="store_true",
                        help="add Postman folders for blueprints")
    args = parser.parse_args()

    logging.disable(logging.CRITICAL)

    try:
        app_path, app_name = args.flask_instance.rsplit('.', 1)
        app = getattr(__import__(app_path), app_name)
    except Exception as e:
        parser.error("can't import \"{}\" ({})".format(args.flask_instance, str(e)))

    if not isinstance(app, Flask):
        parser.error("\"{}\" is not a Flask instance".format(args.flask_instance))

    with app.app_context():
        collection = Collection(args.name)
        rules = list(current_app.url_map.iter_rules())
        for rule in rules:
            for method in rule.methods:
                if args.all or method not in ["OPTIONS", "HEAD"]:
                    endpoint = current_app.view_functions[rule.endpoint]
                    route = Route.from_werkzeug(rule, method, args.base_url)
                    route.description = trim(endpoint.__doc__)
                    folder_name = rule.endpoint.rsplit('.', maxsplit=2)
                    if len(folder_name) >= 2 and args.folders:
                        folder = Folder.get_or_create_folder(collection, folder_name[0])
                        folder.add_route(route)
                    else:
                        collection.add_route(route)

    if args.indent:
        json = json.dumps(collection.to_dict(), indent=4, sort_keys=True)
    else:
        json = json.dumps(collection.to_dict())

    print(json)

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before>from UserDict import IterableUserDict
from contextlib import contextmanager
from datetime import datetime
from uuid import uuid4

from werkzeug import cached_property, LocalProxy
from ZODB.DB import DB
from flask import _request_ctx_stack, current_app
import transaction
from persistent import Persistent
from persistent.list import PersistentList
from persistent.mapping import PersistentMapping
from BTrees.OOBTree import OOBTree

__all__ = ('ZODB', 'Model', 'Factory',
           'List', 'Mapping', 'Timestamp', 'UUID4', 'current_db')

class ZODB(IterableUserDict):
    """ZODB extension for Flask: persistence of native Python objects.

    Basic setup::

        from ZODB.FileStorage import FileStorage
        from flask import Flask, redirect, render_template_string
        from flaskext.zodb import ZODB

        ZODB_STORAGE = lambda: FileStorage('app.fs')

        app = Flask(__name__)
        app.config.from_object(__name__)
        db = ZODB(app)

        @app.route('/')
        @app.route('/<message>')
        def index(message=None):
            if message is not None:
                db['message'] = message
                return redirect('/')
            return render_template_string('Latest message: {{ message }}',
                                          message=db['message'])

    During requests the ``db`` object acts like a Python ``dict``. Any changes
    made to it *directly* will persist, changes made to mutable objects
    within will have to be marked as mutated. This is done for you if you
    inherit :class:`Model` for your own classes and use :attr:`List` and
    :attr:`Mapping` as substitutes for Python's ``list`` and ``dict``.

    Outside of requests, the object can be used as a context manager if
    called, yielding the root object to be used inside the context.
    See :meth:`__call__`.

    """

    def __init__(self, app=None):
        self.app = app
        if app is not None:
            self.init_app(app)

    @cached_property
    def db(self):
        return DB(self.app.config['ZODB_STORAGE']())

    @property
    def connection(self):
        """Request-local database connection."""
        if not hasattr(current_ctx, 'zodb_connection'):
            current_ctx.zodb_connection = self.db.open()
            transaction.begin()
        return current_ctx.zodb_connection

    def get_root(self, root):
        return root

    def root_factory(self, get_root):
        self.get_root = get_root
        return get_root

    @property
    def root(self):
        """Root object for the request-local connection."""
        if not hasattr(current_ctx, 'zodb_root'):
            current_ctx.zodb_root = self.get_root(self.connection.root())
        return current_ctx.zodb_root

    @property
    def data(self):
        return self.connection.root()

    def init_app(self, app):
        assert 'ZODB_STORAGE' in app.config, \
               'ZODB_STORAGE must be configured.'
        self.app = app
        if not hasattr(app, 'extensions'):
            app.extensions = {}
        app.extensions['zodb'] = self

        @app.teardown_request
        def close_db(exception):
            if exception is None:
                transaction.commit()
            else:
                transaction.abort()
            self.connection.close()

    @contextmanager
    def transaction(self):
        connection = self.db.open()
        transaction.begin()
        try:
            yield connection.root()
        except:
            transaction.abort()
        else:
            transaction.commit()
        finally:
            connection.close()

class Factory(object):
    """Set a :class:`Model` attribute with a callable on instantiation.
    Useful for delaying initiation of mutable or dynamic objects.

    ::

        class Dice(Model):
            side = Factory(random.randint, 1, 6)

    ::

        >>> Dice()
        Dice(side=3)
        >>> Dice()
        Dice(side=5)

    """

    def __init__(self, callable, *args, **kwargs):
        self.callable = callable
        self.args = args
        self.kwargs = kwargs

    def __call__(self):
        return self.callable(*self.args, **self.kwargs)

#: UTC timestamp factory
Timestamp = Factory(datetime.utcnow)

#: UUID4 factory
UUID4 = Factory(uuid4)

#: Factory for :func:`PersistentList`
List = Factory(PersistentList)

#: Factory for :func:`PersistentMapping`
Mapping = Factory(PersistentMapping)

#: Factory for an object-to-object balance tree mapping,
#: a :class:`~BTrees.OOBTree.OOBTree`.
BTree = Factory(OOBTree)

class Model(Persistent):
    """Convinience model base.

    You can subclass :class:`persistent.Persistent` directly if you prefer,
    but this base provides some conviniences.

    Set attributes in instantiation::

        >>> Model(title='Hello!')
        Model(title='Hello!')
        >> Model(title='Hello!').title
        'Hello!'

    Declare mutable and dynamic attributes in the class definition::

        class Post(Model):
            id = UUID4
            posted_on = Timestamp
            comments = List

    ::

        >>> Post()
        Post(id=UUID('c3f043a8-8f1f-4381-89b3-fd1f35265925'),
             posted_on=datetime.datetime(2010, 10, 20, 15, 42, 34, 138015),
             comments=[])
        >>> type(Post().comments)
        <class 'persistent.list.PersistentList'>

    """

    def __init__(self, **kwargs):
        for name in dir(self):
            value = getattr(self, name)
            if isinstance(value, Factory):
                setattr(self, name, value())

        for name, value in kwargs.iteritems():
            try:
                attribute = getattr(self, name)
            except AttributeError:
                attribute = None
            if isinstance(attribute, PersistentList):
                attribute.extend(value)
            elif isinstance(attribute, (PersistentMapping, OOBTree)):
                attribute.update(value)
            else:
                setattr(self, name, value)

    def __repr__(self):
        attributes = ', '.join('{0}={1!r}'.format(name, value)
            for (name, value) in vars(self).iteritems())
        return '{0}({1})'.format(self.__class__.__name__, attributes)

#: The :class:`ZODB` instance for the current :class:`~flask.Flask`
#: application.
current_db = LocalProxy(lambda: current_app.extensions['zodb'])

current_ctx = LocalProxy(lambda: _request_ctx_stack.top)
<commit_msg>__all__ is for star-imports, meh<commit_after>from UserDict import IterableUserDict
from contextlib import contextmanager
from datetime import datetime
from uuid import uuid4

from werkzeug import cached_property, LocalProxy
from ZODB.DB import DB
from flask import _request_ctx_stack, current_app
import transaction
from persistent import Persistent
from persistent.list import PersistentList
from persistent.mapping import PersistentMapping
from BTrees.OOBTree import OOBTree

class ZODB(IterableUserDict):
    """ZODB extension for Flask: persistence of native Python objects.

    Basic setup::

        from ZODB.FileStorage import FileStorage
        from flask import Flask, redirect, render_template_string
        from flaskext.zodb import ZODB

        ZODB_STORAGE = lambda: FileStorage('app.fs')

        app = Flask(__name__)
        app.config.from_object(__name__)
        db = ZODB(app)

        @app.route('/')
        @app.route('/<message>')
        def index(message=None):
            if message is not None:
                db['message'] = message
                return redirect('/')
            return render_template_string('Latest message: {{ message }}',
                                          message=db['message'])

    During requests the ``db`` object acts like a Python ``dict``. Any changes
    made to it *directly* will persist, changes made to mutable objects
    within will have to be marked as mutated. This is done for you if you
    inherit :class:`Model` for your own classes and use :attr:`List` and
    :attr:`Mapping` as substitutes for Python's ``list`` and ``dict``.

    Outside of requests, the object can be used as a context manager if
    called, yielding the root object to be used inside the context.
    See :meth:`__call__`.

    """

    def __init__(self, app=None):
        self.app = app
        if app is not None:
            self.init_app(app)

    @cached_property
    def db(self):
        return DB(self.app.config['ZODB_STORAGE']())

    @property
    def connection(self):
        """Request-local database connection."""
        if not hasattr(current_ctx, 'zodb_connection'):
            current_ctx.zodb_connection = self.db.open()
            transaction.begin()
        return current_ctx.zodb_connection

    def get_root(self, root):
        return root

    def root_factory(self, get_root):
        self.get_root = get_root
        return get_root

    @property
    def root(self):
        """Root object for the request-local connection."""
        if not hasattr(current_ctx, 'zodb_root'):
            current_ctx.zodb_root = self.get_root(self.connection.root())
        return current_ctx.zodb_root

    @property
    def data(self):
        return self.connection.root()

    def init_app(self, app):
        assert 'ZODB_STORAGE' in app.config, \
               'ZODB_STORAGE must be configured.'
        self.app = app
        if not hasattr(app, 'extensions'):
            app.extensions = {}
        app.extensions['zodb'] = self

        @app.teardown_request
        def close_db(exception):
            if exception is None:
                transaction.commit()
            else:
                transaction.abort()
            self.connection.close()

    @contextmanager
    def transaction(self):
        connection = self.db.open()
        transaction.begin()
        try:
            yield connection.root()
        except:
            transaction.abort()
        else:
            transaction.commit()
        finally:
            connection.close()

class Factory(object):
    """Set a :class:`Model` attribute with a callable on instantiation.
    Useful for delaying initiation of mutable or dynamic objects.

    ::

        class Dice(Model):
            side = Factory(random.randint, 1, 6)

    ::

        >>> Dice()
        Dice(side=3)
        >>> Dice()
        Dice(side=5)

    """

    def __init__(self, callable, *args, **kwargs):
        self.callable = callable
        self.args = args
        self.kwargs = kwargs

    def __call__(self):
        return self.callable(*self.args, **self.kwargs)

#: UTC timestamp factory
Timestamp = Factory(datetime.utcnow)

#: UUID4 factory
UUID4 = Factory(uuid4)

#: Factory for :func:`PersistentList`
List = Factory(PersistentList)

#: Factory for :func:`PersistentMapping`
Mapping = Factory(PersistentMapping)

#: Factory for an object-to-object balance tree mapping,
#: a :class:`~BTrees.OOBTree.OOBTree`.
BTree = Factory(OOBTree)

class Model(Persistent):
    """Convinience model base.

    You can subclass :class:`persistent.Persistent` directly if you prefer,
    but this base provides some conviniences.

    Set attributes in instantiation::

        >>> Model(title='Hello!')
        Model(title='Hello!')
        >> Model(title='Hello!').title
        'Hello!'

    Declare mutable and dynamic attributes in the class definition::

        class Post(Model):
            id = UUID4
            posted_on = Timestamp
            comments = List

    ::

        >>> Post()
        Post(id=UUID('c3f043a8-8f1f-4381-89b3-fd1f35265925'),
             posted_on=datetime.datetime(2010, 10, 20, 15, 42, 34, 138015),
             comments=[])
        >>> type(Post().comments)
        <class 'persistent.list.PersistentList'>

    """

    def __init__(self, **kwargs):
        for name in dir(self):
            value = getattr(self, name)
            if isinstance(value, Factory):
                setattr(self, name, value())

        for name, value in kwargs.iteritems():
            try:
                attribute = getattr(self, name)
            except AttributeError:
                attribute = None
            if isinstance(attribute, PersistentList):
                attribute.extend(value)
            elif isinstance(attribute, (PersistentMapping, OOBTree)):
                attribute.update(value)
            else:
                setattr(self, name, value)

    def __repr__(self):
        attributes = ', '.join('{0}={1!r}'.format(name, value)
            for (name, value) in vars(self).iteritems())
        return '{0}({1})'.format(self.__class__.__name__, attributes)

#: The :class:`ZODB` instance for the current :class:`~flask.Flask`
#: application.
current_db = LocalProxy(lambda: current_app.extensions['zodb'])

current_ctx = LocalProxy(lambda: _request_ctx_stack.top)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
This is a shim that handles checking and updating salt thin and
then invoking thin.

This is not intended to be instantiated as a module, rather it is a
helper script used by salt.client.ssh.Single.  It is here, in a
separate file, for convenience of development.
'''

import hashlib
import tarfile
import shutil
import sys
import os
import stat

THIN_ARCHIVE = 'salt-thin.tgz'
EXT_ARCHIVE = 'salt-ext_mods.tgz'

# Keep these in sync with salt/exitcodes.py
EX_THIN_DEPLOY = 11
EX_THIN_CHECKSUM = 12
EX_MOD_DEPLOY = 13

class OBJ(object):
    """An empty class for holding instance attribute values."""
    pass

OPTIONS = None
ARGS = None
# The below line is where OPTIONS can be redefined with internal options
# (rather than cli arguments) when the shim is bundled by
# client.ssh.Single._cmd_str()
# pylint: disable=block-comment-should-start-with-'# '
#%%OPTS

def need_deployment():
    """
    Salt thin needs to be deployed - prep the target directory and emit the
    delimeter and exit code that signals a required deployment.
    """
    if os.path.exists(OPTIONS.saltdir):
        shutil.rmtree(OPTIONS.saltdir)
    old_umask = os.umask(0077)
    os.makedirs(OPTIONS.saltdir)
    os.umask(old_umask)
    # Verify perms on saltdir
    euid = os.geteuid()
    dstat = os.stat(OPTIONS.saltdir)
    if dstat.st_uid != euid:
        # Attack detected, try again
        need_deployment()
    if dstat.st_mode != 16832:
        # Attack detected
        need_deployment()
    # If SUDOing then also give the super user group write permissions
    sudo_gid = os.environ.get('SUDO_GID')
    if sudo_gid:
        os.chown(OPTIONS.saltdir, -1, int(sudo_gid))
        stt = os.stat(OPTIONS.saltdir)
        os.chmod(OPTIONS.saltdir, stt.st_mode | stat.S_IWGRP | stat.S_IRGRP | stat.S_IXGRP)

    # Delimeter emitted on stdout *only* to indicate shim message to master.
    sys.stdout.write("{0}\ndeploy\n".format(OPTIONS.delimiter))
    sys.exit(EX_THIN_DEPLOY)

# Adapted from salt.utils.get_hash()
def get_hash(path, form='sha1', chunk_size=4096):
    """Generate a hash digest string for a file."""
    try:
        hash_type = getattr(hashlib, form)
    except AttributeError:
        raise ValueError('Invalid hash type: {0}'.format(form))
    with open(path, 'rb') as ifile:
        hash_obj = hash_type()
        # read the file in in chunks, not the entire file
        for chunk in iter(lambda: ifile.read(chunk_size), b''):
            hash_obj.update(chunk)
        return hash_obj.hexdigest()

def unpack_thin(thin_path):
    """Unpack the Salt thin archive."""
    tfile = tarfile.TarFile.gzopen(thin_path)
    tfile.extractall(path=OPTIONS.saltdir)
    tfile.close()
    os.unlink(thin_path)

def need_ext():
    """Signal that external modules need to be deployed."""
    sys.stdout.write("{0}\next_mods\n".format(OPTIONS.delimiter))
    sys.exit(EX_MOD_DEPLOY)

def unpack_ext(ext_path):
    """Unpack the external modules."""
    modcache = os.path.join(
            OPTIONS.saltdir,
            'running_data',
            'var',
            'cache',
            'salt',
            'minion',
            'extmods')
    tfile = tarfile.TarFile.gzopen(ext_path)
    tfile.extractall(path=modcache)
    tfile.close()
    os.unlink(ext_path)
    ver_path = os.path.join(modcache, 'ext_version')
    ver_dst = os.path.join(OPTIONS.saltdir, 'ext_version')
    shutil.move(ver_path, ver_dst)

def main(argv):  # pylint: disable=W0613
    """Main program body"""
    thin_path = os.path.join(OPTIONS.saltdir, THIN_ARCHIVE)
    if os.path.isfile(thin_path):
        if OPTIONS.checksum != get_hash(thin_path, OPTIONS.hashfunc):
            sys.stderr.write('{0}\n'.format(OPTIONS.checksum))
            sys.stderr.write('{0}\n'.format(get_hash(thin_path, OPTIONS.hashfunc)))
            os.unlink(thin_path)
            sys.stderr.write('WARNING: checksum mismatch for "{0}"\n'.format(thin_path))
            sys.exit(EX_THIN_CHECKSUM)
        unpack_thin(thin_path)
        # Salt thin now is available to use
    else:
        if not os.path.exists(OPTIONS.saltdir):
            need_deployment()

        if not os.path.isdir(OPTIONS.saltdir):
            sys.stderr.write(
                'ERROR: salt path "{0}" exists but is'
                ' not a directory\n'.format(OPTIONS.saltdir)
            )
            sys.exit(os.EX_CANTCREAT)

        version_path = os.path.join(OPTIONS.saltdir, 'version')
        if not os.path.exists(version_path) or not os.path.isfile(version_path):
            sys.stderr.write(
                'WARNING: Unable to locate current thin '
                ' version: {0}.\n'.format(version_path)
            )
            need_deployment()
        with open(version_path, 'r') as vpo:
            cur_version = vpo.readline().strip()
        if cur_version != OPTIONS.version:
            sys.stderr.write(
                'WARNING: current thin version {0}'
                ' is not up-to-date with {1}.\n'.format(
                    cur_version, OPTIONS.version
                )
            )
            need_deployment()
        # Salt thin exists and is up-to-date - fall through and use it

    salt_call_path = os.path.join(OPTIONS.saltdir, 'salt-call')
    if not os.path.isfile(salt_call_path):
        sys.stderr.write('ERROR: thin is missing "{0}"\n'.format(salt_call_path))
        need_deployment()

    with open(os.path.join(OPTIONS.saltdir, 'minion'), 'w') as config:
        config.write(OPTIONS.config + '\n')
    if OPTIONS.ext_mods:
        ext_path = os.path.join(OPTIONS.saltdir, EXT_ARCHIVE)
        if os.path.exists(ext_path):
            unpack_ext(ext_path)
        else:
            version_path = os.path.join(OPTIONS.saltdir, 'ext_version')
            if not os.path.exists(version_path) or not os.path.isfile(version_path):
                need_ext()
            with open(version_path, 'r') as vpo:
                cur_version = vpo.readline().strip()
            if cur_version != OPTIONS.ext_mods:
                need_ext()
    # Fix parameter passing issue
    if len(ARGS) == 1:
        argv_prepared = ARGS[0].split()
    else:
        argv_prepared = ARGS

    salt_argv = [
        sys.executable,
        salt_call_path,
        '--local',
        '--metadata',
        '--out', 'json',
        '-l', 'quiet',
        '-c', OPTIONS.saltdir,
        '--',
    ] + argv_prepared

    sys.stderr.write('SALT_ARGV: {0}\n'.format(salt_argv))

    # Only emit the delimiter on *both* stdout and stderr when completely successful.
    # Yes, the flush() is necessary.
    sys.stdout.write(OPTIONS.delimiter + '\n')
    sys.stdout.flush()
    sys.stderr.write(OPTIONS.delimiter + '\n')
    sys.stderr.flush()
    if OPTIONS.tty:
        import subprocess
        stdout, _ = subprocess.Popen(salt_argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE).communicate()
        sys.stdout.write(stdout)
        sys.stdout.flush()
        if OPTIONS.wipe:
            shutil.rmtree(OPTIONS.saltdir)
    elif OPTIONS.wipe:
        import subprocess
        subprocess.call(salt_argv)
        shutil.rmtree(OPTIONS.saltdir)
    else:
        os.execv(sys.executable, salt_argv)

if __name__ == '__main__':
    sys.exit(main(sys.argv))
<commit_msg>Update to the new disable alias<commit_after># -*- coding: utf-8 -*-
'''
This is a shim that handles checking and updating salt thin and
then invoking thin.

This is not intended to be instantiated as a module, rather it is a
helper script used by salt.client.ssh.Single.  It is here, in a
separate file, for convenience of development.
'''

import hashlib
import tarfile
import shutil
import sys
import os
import stat

THIN_ARCHIVE = 'salt-thin.tgz'
EXT_ARCHIVE = 'salt-ext_mods.tgz'

# Keep these in sync with salt/exitcodes.py
EX_THIN_DEPLOY = 11
EX_THIN_CHECKSUM = 12
EX_MOD_DEPLOY = 13

class OBJ(object):
    """An empty class for holding instance attribute values."""
    pass

OPTIONS = None
ARGS = None
# The below line is where OPTIONS can be redefined with internal options
# (rather than cli arguments) when the shim is bundled by
# client.ssh.Single._cmd_str()
# pylint: disable=block-comment-should-start-with-cardinal-space
#%%OPTS

def need_deployment():
    """
    Salt thin needs to be deployed - prep the target directory and emit the
    delimeter and exit code that signals a required deployment.
    """
    if os.path.exists(OPTIONS.saltdir):
        shutil.rmtree(OPTIONS.saltdir)
    old_umask = os.umask(0077)
    os.makedirs(OPTIONS.saltdir)
    os.umask(old_umask)
    # Verify perms on saltdir
    euid = os.geteuid()
    dstat = os.stat(OPTIONS.saltdir)
    if dstat.st_uid != euid:
        # Attack detected, try again
        need_deployment()
    if dstat.st_mode != 16832:
        # Attack detected
        need_deployment()
    # If SUDOing then also give the super user group write permissions
    sudo_gid = os.environ.get('SUDO_GID')
    if sudo_gid:
        os.chown(OPTIONS.saltdir, -1, int(sudo_gid))
        stt = os.stat(OPTIONS.saltdir)
        os.chmod(OPTIONS.saltdir, stt.st_mode | stat.S_IWGRP | stat.S_IRGRP | stat.S_IXGRP)

    # Delimeter emitted on stdout *only* to indicate shim message to master.
    sys.stdout.write("{0}\ndeploy\n".format(OPTIONS.delimiter))
    sys.exit(EX_THIN_DEPLOY)

# Adapted from salt.utils.get_hash()
def get_hash(path, form='sha1', chunk_size=4096):
    """Generate a hash digest string for a file."""
    try:
        hash_type = getattr(hashlib, form)
    except AttributeError:
        raise ValueError('Invalid hash type: {0}'.format(form))
    with open(path, 'rb') as ifile:
        hash_obj = hash_type()
        # read the file in in chunks, not the entire file
        for chunk in iter(lambda: ifile.read(chunk_size), b''):
            hash_obj.update(chunk)
        return hash_obj.hexdigest()

def unpack_thin(thin_path):
    """Unpack the Salt thin archive."""
    tfile = tarfile.TarFile.gzopen(thin_path)
    tfile.extractall(path=OPTIONS.saltdir)
    tfile.close()
    os.unlink(thin_path)

def need_ext():
    """Signal that external modules need to be deployed."""
    sys.stdout.write("{0}\next_mods\n".format(OPTIONS.delimiter))
    sys.exit(EX_MOD_DEPLOY)

def unpack_ext(ext_path):
    """Unpack the external modules."""
    modcache = os.path.join(
            OPTIONS.saltdir,
            'running_data',
            'var',
            'cache',
            'salt',
            'minion',
            'extmods')
    tfile = tarfile.TarFile.gzopen(ext_path)
    tfile.extractall(path=modcache)
    tfile.close()
    os.unlink(ext_path)
    ver_path = os.path.join(modcache, 'ext_version')
    ver_dst = os.path.join(OPTIONS.saltdir, 'ext_version')
    shutil.move(ver_path, ver_dst)

def main(argv):  # pylint: disable=W0613
    """Main program body"""
    thin_path = os.path.join(OPTIONS.saltdir, THIN_ARCHIVE)
    if os.path.isfile(thin_path):
        if OPTIONS.checksum != get_hash(thin_path, OPTIONS.hashfunc):
            sys.stderr.write('{0}\n'.format(OPTIONS.checksum))
            sys.stderr.write('{0}\n'.format(get_hash(thin_path, OPTIONS.hashfunc)))
            os.unlink(thin_path)
            sys.stderr.write('WARNING: checksum mismatch for "{0}"\n'.format(thin_path))
            sys.exit(EX_THIN_CHECKSUM)
        unpack_thin(thin_path)
        # Salt thin now is available to use
    else:
        if not os.path.exists(OPTIONS.saltdir):
            need_deployment()

        if not os.path.isdir(OPTIONS.saltdir):
            sys.stderr.write(
                'ERROR: salt path "{0}" exists but is'
                ' not a directory\n'.format(OPTIONS.saltdir)
            )
            sys.exit(os.EX_CANTCREAT)

        version_path = os.path.join(OPTIONS.saltdir, 'version')
        if not os.path.exists(version_path) or not os.path.isfile(version_path):
            sys.stderr.write(
                'WARNING: Unable to locate current thin '
                ' version: {0}.\n'.format(version_path)
            )
            need_deployment()
        with open(version_path, 'r') as vpo:
            cur_version = vpo.readline().strip()
        if cur_version != OPTIONS.version:
            sys.stderr.write(
                'WARNING: current thin version {0}'
                ' is not up-to-date with {1}.\n'.format(
                    cur_version, OPTIONS.version
                )
            )
            need_deployment()
        # Salt thin exists and is up-to-date - fall through and use it

    salt_call_path = os.path.join(OPTIONS.saltdir, 'salt-call')
    if not os.path.isfile(salt_call_path):
        sys.stderr.write('ERROR: thin is missing "{0}"\n'.format(salt_call_path))
        need_deployment()

    with open(os.path.join(OPTIONS.saltdir, 'minion'), 'w') as config:
        config.write(OPTIONS.config + '\n')
    if OPTIONS.ext_mods:
        ext_path = os.path.join(OPTIONS.saltdir, EXT_ARCHIVE)
        if os.path.exists(ext_path):
            unpack_ext(ext_path)
        else:
            version_path = os.path.join(OPTIONS.saltdir, 'ext_version')
            if not os.path.exists(version_path) or not os.path.isfile(version_path):
                need_ext()
            with open(version_path, 'r') as vpo:
                cur_version = vpo.readline().strip()
            if cur_version != OPTIONS.ext_mods:
                need_ext()
    # Fix parameter passing issue
    if len(ARGS) == 1:
        argv_prepared = ARGS[0].split()
    else:
        argv_prepared = ARGS

    salt_argv = [
        sys.executable,
        salt_call_path,
        '--local',
        '--metadata',
        '--out', 'json',
        '-l', 'quiet',
        '-c', OPTIONS.saltdir,
        '--',
    ] + argv_prepared

    sys.stderr.write('SALT_ARGV: {0}\n'.format(salt_argv))

    # Only emit the delimiter on *both* stdout and stderr when completely successful.
    # Yes, the flush() is necessary.
    sys.stdout.write(OPTIONS.delimiter + '\n')
    sys.stdout.flush()
    sys.stderr.write(OPTIONS.delimiter + '\n')
    sys.stderr.flush()
    if OPTIONS.tty:
        import subprocess
        stdout, _ = subprocess.Popen(salt_argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE).communicate()
        sys.stdout.write(stdout)
        sys.stdout.flush()
        if OPTIONS.wipe:
            shutil.rmtree(OPTIONS.saltdir)
    elif OPTIONS.wipe:
        import subprocess
        subprocess.call(salt_argv)
        shutil.rmtree(OPTIONS.saltdir)
    else:
        os.execv(sys.executable, salt_argv)

if __name__ == '__main__':
    sys.exit(main(sys.argv))
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
Service support for Debian systems (uses update-rc.d and /sbin/service)
'''

# Import python libs
import glob
import re

# Import salt libs
from .systemd import _sd_booted

__func_alias__ = {
    'reload_': 'reload'
}

# Define the module's virtual name
__virtualname__ = 'service'

def __virtual__():
    '''
    Only work on Debian and when systemd isn't running
    '''
    if __grains__['os'] in ('Debian', 'Raspbian') and not _sd_booted():
        return __virtualname__
    return False

def _get_runlevel():
    '''
    returns the current runlevel
    '''
    out = __salt__['cmd.run']('runlevel')
    # unknown can be returned while inside a container environment, since
    # this is due to a lack of init, it should be safe to assume runlevel
    # 2, which is Debian's default. If not, all service related states
    # will throw an out of range exception here which will cause
    # other functions to fail.
    if 'unknown' in out:
        return '2'
    else:
        return out.split()[1]

def get_enabled():
    '''
    Return a list of service that are enabled on boot

    CLI Example:

    .. code-block:: bash

        salt '*' service.get_enabled
    '''
    prefix = '/etc/rc[S{0}].d/S'.format(_get_runlevel())
    ret = set()
    lines = glob.glob('{0}*'.format(prefix))
    for line in lines:
        ret.add(re.split(prefix + r'\d+', line)[1])
    return sorted(ret)

def get_disabled():
    '''
    Return a set of services that are installed but disabled

    CLI Example:

    .. code-block:: bash

        salt '*' service.get_disabled
    '''
    return sorted(set(get_all()) - set(get_enabled()))

def available(name):
    '''
    Returns ``True`` if the specified service is available, otherwise returns
    ``False``.

    CLI Example:

    .. code-block:: bash

        salt '*' service.available sshd
    '''
    return name in get_all()

def get_all():
    '''
    Return all available boot services

    CLI Example:

    .. code-block:: bash

        salt '*' service.get_all
    '''
    ret = set()
    lines = glob.glob('/etc/init.d/*')
    for line in lines:
        service = line.split('/etc/init.d/')[1]
        # Remove README.  If it's an enabled service, it will be added back in.
        if service != 'README':
            ret.add(service)
    return sorted(ret | set(get_enabled()))

def start(name):
    '''
    Start the specified service

    CLI Example:

    .. code-block:: bash

        salt '*' service.start <service name>
    '''
    cmd = 'service {0} start'.format(name)
    return not __salt__['cmd.retcode'](cmd)

def stop(name):
    '''
    Stop the specified service

    CLI Example:

    .. code-block:: bash

        salt '*' service.stop <service name>
    '''
    cmd = 'service {0} stop'.format(name)
    return not __salt__['cmd.retcode'](cmd)

def restart(name):
    '''
    Restart the named service

    CLI Example:

    .. code-block:: bash

        salt '*' service.restart <service name>
    '''
    cmd = 'service {0} restart'.format(name)
    return not __salt__['cmd.retcode'](cmd)

def reload_(name):
    '''
    Reload the named service

    CLI Example:

    .. code-block:: bash

        salt '*' service.reload <service name>
    '''
    cmd = 'service {0} reload'.format(name)
    return not __salt__['cmd.retcode'](cmd)

def force_reload(name):
    '''
    Force-reload the named service

    CLI Example:

    .. code-block:: bash

        salt '*' service.force_reload <service name>
    '''
    cmd = 'service {0} force-reload'.format(name)
    return not __salt__['cmd.retcode'](cmd)

def status(name, sig=None):
    '''
    Return the status for a service, pass a signature to use to find
    the service via ps

    CLI Example:

    .. code-block:: bash

        salt '*' service.status <service name>
    '''
    if sig:
        return bool(__salt__['status.pid'](sig))
    cmd = 'service {0} status'.format(name)
    return not __salt__['cmd.retcode'](cmd)

def enable(name, **kwargs):
    '''
    Enable the named service to start at boot

    CLI Example:

    .. code-block:: bash

        salt '*' service.enable <service name>
    '''
    cmd = 'update-rc.d {0} enable'.format(name)
    osmajor = __grains__['osrelease'].split('.')[0]
    if int(osmajor) >= 6:
        cmd = 'insserv {0} && '.format(name) + cmd
    return not __salt__['cmd.retcode'](cmd)

def disable(name, **kwargs):
    '''
    Disable the named service to start at boot

    CLI Example:

    .. code-block:: bash

        salt '*' service.disable <service name>
    '''
    cmd = 'update-rc.d {0} disable'.format(name)
    return not __salt__['cmd.retcode'](cmd)

def enabled(name):
    '''
    Return True if the named service is enabled, false otherwise

    CLI Example:

    .. code-block:: bash

        salt '*' service.enabled <service name>
    '''
    return name in get_enabled()

def disabled(name):
    '''
    Return True if the named service is enabled, false otherwise

    CLI Example:

    .. code-block:: bash

        salt '*' service.disabled <service name>
    '''
    return name in get_disabled()
<commit_msg>Added missing function to debian<commit_after># -*- coding: utf-8 -*-
'''
Service support for Debian systems (uses update-rc.d and /sbin/service)
'''

# Import python libs
import glob
import re

# Import salt libs
from .systemd import _sd_booted

__func_alias__ = {
    'reload_': 'reload'
}

# Define the module's virtual name
__virtualname__ = 'service'

def __virtual__():
    '''
    Only work on Debian and when systemd isn't running
    '''
    if __grains__['os'] in ('Debian', 'Raspbian') and not _sd_booted():
        return __virtualname__
    return False

def _get_runlevel():
    '''
    returns the current runlevel
    '''
    out = __salt__['cmd.run']('runlevel')
    # unknown can be returned while inside a container environment, since
    # this is due to a lack of init, it should be safe to assume runlevel
    # 2, which is Debian's default. If not, all service related states
    # will throw an out of range exception here which will cause
    # other functions to fail.
    if 'unknown' in out:
        return '2'
    else:
        return out.split()[1]

def get_enabled():
    '''
    Return a list of service that are enabled on boot

    CLI Example:

    .. code-block:: bash

        salt '*' service.get_enabled
    '''
    prefix = '/etc/rc[S{0}].d/S'.format(_get_runlevel())
    ret = set()
    lines = glob.glob('{0}*'.format(prefix))
    for line in lines:
        ret.add(re.split(prefix + r'\d+', line)[1])
    return sorted(ret)

def get_disabled():
    '''
    Return a set of services that are installed but disabled

    CLI Example:

    .. code-block:: bash

        salt '*' service.get_disabled
    '''
    return sorted(set(get_all()) - set(get_enabled()))

def available(name):
    '''
    Returns ``True`` if the specified service is available, otherwise returns
    ``False``.

    CLI Example:

    .. code-block:: bash

        salt '*' service.available sshd
    '''
    return name in get_all()

def missing(name):
    '''
    The inverse of service.available.
    Returns ``True`` if the specified service is not available, otherwise returns
    ``False``.

    CLI Example:

    .. code-block:: bash

        salt '*' service.missing sshd
    '''
    return not name in get_all()

def get_all():
    '''
    Return all available boot services

    CLI Example:

    .. code-block:: bash

        salt '*' service.get_all
    '''
    ret = set()
    lines = glob.glob('/etc/init.d/*')
    for line in lines:
        service = line.split('/etc/init.d/')[1]
        # Remove README.  If it's an enabled service, it will be added back in.
        if service != 'README':
            ret.add(service)
    return sorted(ret | set(get_enabled()))

def start(name):
    '''
    Start the specified service

    CLI Example:

    .. code-block:: bash

        salt '*' service.start <service name>
    '''
    cmd = 'service {0} start'.format(name)
    return not __salt__['cmd.retcode'](cmd)

def stop(name):
    '''
    Stop the specified service

    CLI Example:

    .. code-block:: bash

        salt '*' service.stop <service name>
    '''
    cmd = 'service {0} stop'.format(name)
    return not __salt__['cmd.retcode'](cmd)

def restart(name):
    '''
    Restart the named service

    CLI Example:

    .. code-block:: bash

        salt '*' service.restart <service name>
    '''
    cmd = 'service {0} restart'.format(name)
    return not __salt__['cmd.retcode'](cmd)

def reload_(name):
    '''
    Reload the named service

    CLI Example:

    .. code-block:: bash

        salt '*' service.reload <service name>
    '''
    cmd = 'service {0} reload'.format(name)
    return not __salt__['cmd.retcode'](cmd)

def force_reload(name):
    '''
    Force-reload the named service

    CLI Example:

    .. code-block:: bash

        salt '*' service.force_reload <service name>
    '''
    cmd = 'service {0} force-reload'.format(name)
    return not __salt__['cmd.retcode'](cmd)

def status(name, sig=None):
    '''
    Return the status for a service, pass a signature to use to find
    the service via ps

    CLI Example:

    .. code-block:: bash

        salt '*' service.status <service name>
    '''
    if sig:
        return bool(__salt__['status.pid'](sig))
    cmd = 'service {0} status'.format(name)
    return not __salt__['cmd.retcode'](cmd)

def enable(name, **kwargs):
    '''
    Enable the named service to start at boot

    CLI Example:

    .. code-block:: bash

        salt '*' service.enable <service name>
    '''
    cmd = 'update-rc.d {0} enable'.format(name)
    osmajor = __grains__['osrelease'].split('.')[0]
    if int(osmajor) >= 6:
        cmd = 'insserv {0} && '.format(name) + cmd
    return not __salt__['cmd.retcode'](cmd)

def disable(name, **kwargs):
    '''
    Disable the named service to start at boot

    CLI Example:

    .. code-block:: bash

        salt '*' service.disable <service name>
    '''
    cmd = 'update-rc.d {0} disable'.format(name)
    return not __salt__['cmd.retcode'](cmd)

def enabled(name):
    '''
    Return True if the named service is enabled, false otherwise

    CLI Example:

    .. code-block:: bash

        salt '*' service.enabled <service name>
    '''
    return name in get_enabled()

def disabled(name):
    '''
    Return True if the named service is enabled, false otherwise

    CLI Example:

    .. code-block:: bash

        salt '*' service.disabled <service name>
    '''
    return name in get_disabled()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
Module for viewing and modifying sysctl parameters
'''

# Import Python libs
from __future__ import absolute_import, unicode_literals, print_function
import logging

# Import salt libs
import salt.utils.files
from salt.exceptions import CommandExecutionError
from salt.ext import six

# Define the module's virtual name
__virtualname__ = 'sysctl'

# Get logging started
log = logging.getLogger(__name__)

def __virtual__():
    '''
    Only runs on FreeBSD systems
    '''
    if __grains__['os'] == 'FreeBSD':
        return __virtualname__
    return (False, 'The freebsd_sysctl execution module cannot be loaded: '
            'only available on FreeBSD systems.')

def _formatfor(name, value, config, tail=''):
    if config == '/boot/loader.conf':
        return '{0}="{1}"{2}'.format(name, value, tail)
    else:
        return '{0}={1}{2}'.format(name, value, tail)

def show(config_file=False):
    '''
    Return a list of sysctl parameters for this minion

    CLI Example:

    .. code-block:: bash

        salt '*' sysctl.show
    '''
    roots = (
        'compat',
        'debug',
        'dev',
        'hptmv',
        'hw',
        'kern',
        'machdep',
        'net',
        'p1003_1b',
        'security',
        'user',
        'vfs',
        'vm'
    )
    cmd = 'sysctl -ae'
    ret = {}
    comps = ['']

    if config_file:
        try:
            with salt.utils.files.fopen(config_file, 'r') as f:
                for line in f.readlines():
                    l = line.strip()
                    if l != "" and not l.startswith("#"):
                        comps = line.split('=', 1)
                        ret[comps[0]] = comps[1]
            return ret
        except (OSError, IOError):
            log.error('Could not open sysctl config file')
            return None
    else:
        out = __salt__['cmd.run'](cmd, output_loglevel='trace')
        for line in out.splitlines():
            if any([line.startswith('{0}.'.format(root)) for root in roots]):
                comps = line.split('=', 1)
                ret[comps[0]] = comps[1]
            elif comps[0]:
                ret[comps[0]] += '{0}\n'.format(line)
            else:
                continue
        return ret

def get(name):
    '''
    Return a single sysctl parameter for this minion

    CLI Example:

    .. code-block:: bash

        salt '*' sysctl.get hw.physmem
    '''
    cmd = 'sysctl -n {0}'.format(name)
    out = __salt__['cmd.run'](cmd, python_shell=False)
    return out

def assign(name, value):
    '''
    Assign a single sysctl parameter for this minion

    CLI Example:

    .. code-block:: bash

        salt '*' sysctl.assign net.inet.icmp.icmplim 50
    '''
    ret = {}
    cmd = 'sysctl {0}="{1}"'.format(name, value)
    data = __salt__['cmd.run_all'](cmd, python_shell=False)

    if data['retcode'] != 0:
        raise CommandExecutionError('sysctl failed: {0}'.format(
            data['stderr']))
    new_name, new_value = data['stdout'].split(':', 1)
    ret[new_name] = new_value.split(' -> ')[-1]
    return ret

def persist(name, value, config='/etc/sysctl.conf'):
    '''
    Assign and persist a simple sysctl parameter for this minion

    CLI Example:

    .. code-block:: bash

        salt '*' sysctl.persist net.inet.icmp.icmplim 50
        salt '*' sysctl.persist coretemp_load NO config=/boot/loader.conf
    '''
    nlines = []
    edited = False
    value = six.text_type(value)

    with salt.utils.files.fopen(config, 'r') as ifile:
        for line in ifile:
            line = salt.utils.stringutils.to_unicode(line).rstrip('\n')
            if not line.startswith('{0}='.format(name)):
                nlines.append(line)
                continue
            else:
                key, rest = line.split('=', 1)
                if rest.startswith('"'):
                    _, rest_v, rest = rest.split('"', 2)
                elif rest.startswith('\''):
                    _, rest_v, rest = rest.split('\'', 2)
                else:
                    rest_v = rest.split()[0]
                    rest = rest[len(rest_v):]
                if rest_v == value:
                    return 'Already set'
                new_line = _formatfor(key, value, config, rest)
                nlines.append(new_line)
                edited = True
    if not edited:
        nlines.append("{0}\n".format(_formatfor(name, value, config)))
    with salt.utils.files.fopen(config, 'w+') as ofile:
        nlines = [salt.utils.stringutils.to_str(_l) for _l in nlines]
        ofile.writelines(nlines)
    if config != '/boot/loader.conf':
        assign(name, value)
    return 'Updated'
<commit_msg>Backporting #48087 to 2018.3<commit_after># -*- coding: utf-8 -*-
'''
Module for viewing and modifying sysctl parameters
'''

# Import Python libs
from __future__ import absolute_import, unicode_literals, print_function
import logging

# Import salt libs
import salt.utils.files
from salt.exceptions import CommandExecutionError
from salt.ext import six

# Define the module's virtual name
__virtualname__ = 'sysctl'

# Get logging started
log = logging.getLogger(__name__)

def __virtual__():
    '''
    Only runs on FreeBSD systems
    '''
    if __grains__['os'] == 'FreeBSD':
        return __virtualname__
    return (False, 'The freebsd_sysctl execution module cannot be loaded: '
            'only available on FreeBSD systems.')

def _formatfor(name, value, config, tail=''):
    if config == '/boot/loader.conf':
        return '{0}="{1}"{2}'.format(name, value, tail)
    else:
        return '{0}={1}{2}'.format(name, value, tail)

def show(config_file=False):
    '''
    Return a list of sysctl parameters for this minion

    CLI Example:

    .. code-block:: bash

        salt '*' sysctl.show
    '''
    roots = (
        'compat',
        'debug',
        'dev',
        'hptmv',
        'hw',
        'kern',
        'machdep',
        'net',
        'p1003_1b',
        'security',
        'user',
        'vfs',
        'vm'
    )
    cmd = 'sysctl -ae'
    ret = {}
    comps = ['']

    if config_file:
        try:
            with salt.utils.files.fopen(config_file, 'r') as f:
                for line in f.readlines():
                    l = line.strip()
                    if l != "" and not l.startswith("#"):
                        comps = line.split('=', 1)
                        ret[comps[0]] = comps[1]
            return ret
        except (OSError, IOError):
            log.error('Could not open sysctl config file')
            return None
    else:
        out = __salt__['cmd.run'](cmd, output_loglevel='trace')
        for line in out.splitlines():
            if any([line.startswith('{0}.'.format(root)) for root in roots]):
                comps = line.split('=', 1)
                ret[comps[0]] = comps[1]
            elif comps[0]:
                ret[comps[0]] += '{0}\n'.format(line)
            else:
                continue
        return ret

def get(name):
    '''
    Return a single sysctl parameter for this minion

    CLI Example:

    .. code-block:: bash

        salt '*' sysctl.get hw.physmem
    '''
    cmd = 'sysctl -n {0}'.format(name)
    out = __salt__['cmd.run'](cmd, python_shell=False)
    return out

def assign(name, value):
    '''
    Assign a single sysctl parameter for this minion

    CLI Example:

    .. code-block:: bash

        salt '*' sysctl.assign net.inet.icmp.icmplim 50
    '''
    ret = {}
    cmd = 'sysctl {0}="{1}"'.format(name, value)
    data = __salt__['cmd.run_all'](cmd, python_shell=False)

    if data['retcode'] != 0:
        raise CommandExecutionError('sysctl failed: {0}'.format(
            data['stderr']))
    new_name, new_value = data['stdout'].split(':', 1)
    ret[new_name] = new_value.split(' -> ')[-1]
    return ret

def persist(name, value, config='/etc/sysctl.conf'):
    '''
    Assign and persist a simple sysctl parameter for this minion

    CLI Example:

    .. code-block:: bash

        salt '*' sysctl.persist net.inet.icmp.icmplim 50
        salt '*' sysctl.persist coretemp_load NO config=/boot/loader.conf
    '''
    nlines = []
    edited = False
    value = six.text_type(value)

    with salt.utils.files.fopen(config, 'r') as ifile:
        for line in ifile:
            line = salt.utils.stringutils.to_unicode(line).rstrip('\n')
            if not line.startswith('{0}='.format(name)):
                nlines.append(line)
                continue
            else:
                key, rest = line.split('=', 1)
                if rest.startswith('"'):
                    _, rest_v, rest = rest.split('"', 2)
                elif rest.startswith('\''):
                    _, rest_v, rest = rest.split('\'', 2)
                else:
                    rest_v = rest.split()[0]
                    rest = rest[len(rest_v):]
                if rest_v == value:
                    return 'Already set'
                new_line = _formatfor(key, value, config, rest)
                nlines.append(new_line)
                edited = True
    if not edited:
        nlines.append("{0}\n".format(_formatfor(name, value, config)))
    with salt.utils.files.fopen(config, 'w+') as ofile:
        nlines = [salt.utils.stringutils.to_str(_l) + '\n' for _l in nlines]
        ofile.writelines(nlines)
    if config != '/boot/loader.conf':
        assign(name, value)
    return 'Updated'
<|endoftext|>