<commit_before>from flask import Flask, g, redirect, render_template, request, session
from flask import flash, url_for
from flask_github import GitHub
import os
import requests
from sqlalchemy import create_engine, Column, Integer, String, Boolean
from sqlalchemy.orm import scoped_session, sessionmaker
from sqlalchemy.ext.declarative import declarative_base

app = Flask(__name__)
app.config['GITHUB_CLIENT_ID'] = os.environ['AELITA_GITHUB_CLIENT_ID']
app.config['GITHUB_CLIENT_SECRET'] = os.environ['AELITA_GITHUB_CLIENT_SECRET']
app.config['BOT_USERNAME'] = os.environ['AELITA_BOT_USERNAME']
app.config['BOT_BASEURL'] = os.environ['AELITA_BOT_BASEURL']
app.config['BOT_DBURI'] = os.environ['AELITA_BOT_DBURI']
app.config['BOT_ACCESS_TOKEN'] = os.environ['AELITA_BOT_ACCESS_TOKEN']
app.config['BOT_NOTICE_SECRET'] = os.environ['AELITA_BOT_NOTICE_SECRET']
app.config['BOT_STATUS_SECRET'] = os.environ['AELITA_BOT_STATUS_SECRET']
app.secret_key = os.environ['AELITA_VIEW_SECRET']
github = GitHub(app)
engine = create_engine(app.config['BOT_DBURI'])
db_session = scoped_session(
    sessionmaker(
        autocommit=False,
        autoflush=False,
        bind=engine,
    )
)

Base = declarative_base()
Base.query = db_session.query_property()

class User(Base):
    __tablename__ = 'signup_users'

    user_id = Column(Integer, primary_key=True)
    username = Column(String(200))
    github_access_token = Column(String(200))
    invite_count = Column(Integer, default=3)

    def __init__(self, github_access_token):
        self.github_access_token = github_access_token

class Invited(Base):
    __tablename__ = 'signup_invited'

    username = Column(String(200), primary_key=True)

    def __init__(self, username):
        self.username = username

class Pipeline(Base):
    __tablename__ = 'twelvef_config_pipeline'

    pipeline_id = Column(Integer, primary_key=True)
    name = Column(String(200))

    def __init__(self, pipeline_id, name):
        self.pipeline_id = pipeline_id
        self.name = name

class GithubProjects(Base):
    __tablename__ = 'twelvef_github_projects'

    pipeline_id = Column(Integer, primary_key=True)
    try_pipeline_id = Column(Integer, nullable=True)
    owner = Column(String(200))
    repo = Column(String(200))

    def __init__(self, pipeline_id, try_pipeline_id, owner, repo):
        self.pipeline_id = pipeline_id
        self.try_pipeline_id = try_pipeline_id
        self.owner = owner
        self.repo = repo

class GithubStatusPipelines(Base):
    __tablename__ = 'twelvef_github_status_pipelines'

    pipeline_id = Column(Integer, primary_key=True)
    owner = Column(String(200))
    repo = Column(String(200))
    context = Column(String(200))

    def __init__(self, pipeline_id, owner, repo, context):
        self.pipeline_id = pipeline_id
        self.owner = owner
        self.repo = repo
        self.context = context

class GithubGitPipelines(Base):
    __tablename__ = 'twelvef_github_git_pipelines'

    pipeline_id = Column(Integer, primary_key=True)
    owner = Column(String(200))
    repo = Column(String(200))
    master_branch = Column(String(200))
    staging_branch = Column(String(200))
    push_to_master = Column(Boolean)

    def __init__(self, pipeline_id, owner, repo):
        self.pipeline_id = pipeline_id
        self.owner = owner
        self.repo = repo
        self.master_branch = "master"
        self.staging_branch = "staging"
        self.push_to_master = True

Base.metadata.create_all(bind=engine)

@app.before_request
def before_request():
    g.user = None

@app.after_request
def after_request(response):
    db_session.remove()
    g.user = None
    return response

def get_user():
    if g.user is None and 'user_id' in session:
        g.user = User.query.get(session['user_id'])
    return g.user

@github.access_token_getter
def token_getter():
    user = get_user()
    if user is not None:
        return user.github_access_token

@app.route('/')
def index():
    if 'user_id' not in session:
        return render_template('index.html')
    else:
        return redirect(url_for('manage'))

@app.route('/login')
def login():
    return github.authorize(scope="user,repo")

@app.route('/logout', methods=['POST'])
def logout():
    session.pop('user_id', None)
    return redirect(url_for('index'))

@app.route('/github-callback')
@github.authorized_handler
def authorized(oauth_token):
    if oauth_token is None:
        flash("Authorization failed.")
        return redirect(url_for('index'))
    g.user = User(oauth_token)
    username = github.get('user')['login']
    g.user.username = username
    user = User.query.filter_by(username=username).first()
    if user is None:
        invite = Invited.query.filter_by(username=username).first()
        if invite is None:
            flash("This service is invite-only")
            return redirect(url_for('manage'))
        user = db_session.merge(g.user)
        g.user = user
        db_session.add(user)
        db_session.commit()
    session['user_id'] = user.user_id
    return redirect(url_for('manage'))

@app.route('/manage', methods=['GET', 'POST'])
def manage():
    user = get_user()
    if user is None:
        flash("Please log in")
        return logout()
    all_repos = github.get('user/repos')
    present = []
    non_present = []
    edit = None
    default_context = 'continuous-integration/travis-ci/push'
    for repo in all_repos:
        on_repo = GithubProjects.query \
            .filter_by(owner=repo['owner']['login'],repo=repo['name']) \
            .first()
        repo_def = {
            "owner": repo['owner']['login'],
            "repo": repo['name'],
            "name": repo['full_name'],
            "id": repo['id'],
        }
        if request.method == 'POST':
            if 'add' in request.form and \
                    int(request.form['add']) == repo['id'] and \
                    on_repo is None:
                return add_repo(repo, default_context)
            elif 'remove' in request.form and \
                    int(request.form['remove']) == repo['id'] and \
                    on_repo is not None:
                return remove_repo(repo, on_repo)
            elif 'edit' in request.form and \
                    int(request.form['edit']) == repo['id'] and \
                    on_repo is not None:
                return edit_repo(on_repo)
        if on_repo is None:
            non_present.append(repo_def)
        else:
            present.append(repo_def)
            if request.args.get('edit') and \
                    request.args.get('edit') == str(repo['id']):
                edit = repo_def
                on_status = GithubStatusPipelines.query \
                    .filter_by(pipeline_id=on_repo.pipeline_id) \
                    .first()
                on_git = GithubGitPipelines.query \
                    .filter_by(pipeline_id=on_repo.pipeline_id) \
                    .first()
                edit['context'] = on_status.context
                edit['master_branch'] = on_git.master_branch
                edit['staging_branch'] = on_git.staging_branch
    return render_template(
        'manage.html',
        username=user.username,
        non_present=non_present,
        present=present,
        edit=edit,
        invite_count=user.invite_count,
        base_url=app.config['BOT_BASEURL'],
    )

def add_repo(repo, context):
    user = get_user()
    # Add repository to our database
    pipeline = Pipeline(None, repo['full_name'])
    db_session.add(pipeline)
    db_session.flush()
    pipeline_id = pipeline.pipeline_id
    on_repo = GithubProjects(pipeline_id, None, repo['owner']['login'], repo['name'])
    db_session.add(on_repo)
    status = GithubStatusPipelines(
        pipeline_id,
        repo['owner']['login'],
        repo['name'],
        context
    )
    db_session.add(status)
    git = GithubGitPipelines(
        pipeline_id,
        repo['owner']['login'],
        repo['name']
    )
    db_session.add(git)
    db_session.commit()
    # Add our account as a collaborator on Github
    headers = {
        "Content-Type": "application/vnd.github.swamp-thing-preview+json",
    }
    invite=github.request(
        'PUT',
        'repos/' + repo['full_name'] + '/collaborators/' + \
            app.config['BOT_USERNAME'],
        headers=headers,
    )
    headers = {
        "Accept": "token " + app.config['BOT_ACCESS_TOKEN'],
        "Content-Type": "application/vnd.github.swamp-thing-preview+json",
    }
    requests.request(
        'PATCH',
        invite.url,
        headers=headers,
    )
    # Add our webhooks
    github.post(
        'repos/' + repo['full_name'] + '/hooks',
        {
            "name": "web",
            "active": True,
            "config": {
                "url": app.config['BOT_BASEURL'] + '/github-notice',
                "content_type": "json",
                "secret": app.config['BOT_NOTICE_SECRET']
            },
            "events": [ "issue_comment", "pull_request", "team_add" ]
        }
    )
    github.post(
        'repos/' + repo['full_name'] + '/hooks',
        {
            "name": "web",
            "active": True,
            "config": {
                "url": app.config['BOT_BASEURL'] + '/github-status',
                "content_type": "json",
                "secret": app.config['BOT_STATUS_SECRET']
            },
            "events": [ "status" ]
        }
    )
    flash("Added successfully")
    return redirect(url_for('manage'))

def remove_repo(repo, on_repo):
    user = get_user()
    # Remove from our database
    pipeline_id = on_repo.pipeline_id
    db_session.delete(on_repo)
    pipeline = Pipeline.query.filter_by(pipeline_id=pipeline_id).first()
    if not pipeline is None: db_session.delete(pipeline)
    status = GithubStatusPipelines.query \
            .filter_by(pipeline_id=pipeline_id) \
            .first();
    if not status is None: db_session.delete(status)
    git = GithubGitPipelines.query \
            .filter_by(pipeline_id=pipeline_id) \
            .first();
    if not git is None: db_session.delete(git)
    db_session.commit()
    # Remove our collaboratorship
    github.raw_request(
        'DELETE',
        'repos/' + repo['full_name'] + '/collaborators/' + \
            app.config['BOT_USERNAME']
    )
    # Remove our webhooks
    hooks = github.get(
        'repos/' + repo['full_name'] + '/hooks'
    )
    my_urls = {
        app.config['BOT_BASEURL'] + '/github-notice',
        app.config['BOT_BASEURL'] + '/github-status',
    }
    for webhook in hooks:
        if webhook['name'] == 'web' and webhook['config']['url'] in my_urls:
            github.request(
                'DELETE',
                'repos/' + repo['full_name'] + '/hooks/' + str(webhook['id'])
            )
    flash("Deleted successfully")
    return redirect(url_for('manage'))

def edit_repo(project):
    user = get_user()
    # Remove from our database
    pipeline_id = project.pipeline_id
    status = GithubStatusPipelines.query \
            .filter_by(pipeline_id=pipeline_id) \
            .first();
    status.context = request.form['context']
    git = GithubGitPipelines.query \
            .filter_by(pipeline_id=pipeline_id) \
            .first();
    git.master_branch = request.form['master_branch']
    git.staging_branch = request.form['staging_branch']
    db_session.commit()
    flash("Saved successfully")
    return redirect(url_for('manage'))

@app.route('/invite', methods=['POST'])
def invite():
    user = get_user()
    if user is None:
        flash("Please log in")
        return redirect(url_for('index'))
    if user.invite_count <= 0:
        flash("You're out of invites")
        return redirect(url_for('manage'))
    invite = Invited.query.filter_by(username=request.form['username'])
    if invite.first() is not None:
        flash("This person is already invited")
        return redirect(url_for('manage'))
    invite = Invited(request.form['username'])
    db_session.add(invite)
    user.invite_count = user.invite_count - 1
    db_session.commit()
    flash("Invitation recorded; now let them know!")
    return redirect(url_for('manage'))
<commit_msg>Fix line length<commit_after>from flask import Flask, g, redirect, render_template, request, session
from flask import flash, url_for
from flask_github import GitHub
import os
import requests
from sqlalchemy import create_engine, Column, Integer, String, Boolean
from sqlalchemy.orm import scoped_session, sessionmaker
from sqlalchemy.ext.declarative import declarative_base

app = Flask(__name__)
app.config['GITHUB_CLIENT_ID'] = os.environ['AELITA_GITHUB_CLIENT_ID']
app.config['GITHUB_CLIENT_SECRET'] = os.environ['AELITA_GITHUB_CLIENT_SECRET']
app.config['BOT_USERNAME'] = os.environ['AELITA_BOT_USERNAME']
app.config['BOT_BASEURL'] = os.environ['AELITA_BOT_BASEURL']
app.config['BOT_DBURI'] = os.environ['AELITA_BOT_DBURI']
app.config['BOT_ACCESS_TOKEN'] = os.environ['AELITA_BOT_ACCESS_TOKEN']
app.config['BOT_NOTICE_SECRET'] = os.environ['AELITA_BOT_NOTICE_SECRET']
app.config['BOT_STATUS_SECRET'] = os.environ['AELITA_BOT_STATUS_SECRET']
app.secret_key = os.environ['AELITA_VIEW_SECRET']
github = GitHub(app)
engine = create_engine(app.config['BOT_DBURI'])
db_session = scoped_session(
    sessionmaker(
        autocommit=False,
        autoflush=False,
        bind=engine,
    )
)

Base = declarative_base()
Base.query = db_session.query_property()

class User(Base):
    __tablename__ = 'signup_users'

    user_id = Column(Integer, primary_key=True)
    username = Column(String(200))
    github_access_token = Column(String(200))
    invite_count = Column(Integer, default=3)

    def __init__(self, github_access_token):
        self.github_access_token = github_access_token

class Invited(Base):
    __tablename__ = 'signup_invited'

    username = Column(String(200), primary_key=True)

    def __init__(self, username):
        self.username = username

class Pipeline(Base):
    __tablename__ = 'twelvef_config_pipeline'

    pipeline_id = Column(Integer, primary_key=True)
    name = Column(String(200))

    def __init__(self, pipeline_id, name):
        self.pipeline_id = pipeline_id
        self.name = name

class GithubProjects(Base):
    __tablename__ = 'twelvef_github_projects'

    pipeline_id = Column(Integer, primary_key=True)
    try_pipeline_id = Column(Integer, nullable=True)
    owner = Column(String(200))
    repo = Column(String(200))

    def __init__(self, pipeline_id, try_pipeline_id, owner, repo):
        self.pipeline_id = pipeline_id
        self.try_pipeline_id = try_pipeline_id
        self.owner = owner
        self.repo = repo

class GithubStatusPipelines(Base):
    __tablename__ = 'twelvef_github_status_pipelines'

    pipeline_id = Column(Integer, primary_key=True)
    owner = Column(String(200))
    repo = Column(String(200))
    context = Column(String(200))

    def __init__(self, pipeline_id, owner, repo, context):
        self.pipeline_id = pipeline_id
        self.owner = owner
        self.repo = repo
        self.context = context

class GithubGitPipelines(Base):
    __tablename__ = 'twelvef_github_git_pipelines'

    pipeline_id = Column(Integer, primary_key=True)
    owner = Column(String(200))
    repo = Column(String(200))
    master_branch = Column(String(200))
    staging_branch = Column(String(200))
    push_to_master = Column(Boolean)

    def __init__(self, pipeline_id, owner, repo):
        self.pipeline_id = pipeline_id
        self.owner = owner
        self.repo = repo
        self.master_branch = "master"
        self.staging_branch = "staging"
        self.push_to_master = True

Base.metadata.create_all(bind=engine)

@app.before_request
def before_request():
    g.user = None

@app.after_request
def after_request(response):
    db_session.remove()
    g.user = None
    return response

def get_user():
    if g.user is None and 'user_id' in session:
        g.user = User.query.get(session['user_id'])
    return g.user

@github.access_token_getter
def token_getter():
    user = get_user()
    if user is not None:
        return user.github_access_token

@app.route('/')
def index():
    if 'user_id' not in session:
        return render_template('index.html')
    else:
        return redirect(url_for('manage'))

@app.route('/login')
def login():
    return github.authorize(scope="user,repo")

@app.route('/logout', methods=['POST'])
def logout():
    session.pop('user_id', None)
    return redirect(url_for('index'))

@app.route('/github-callback')
@github.authorized_handler
def authorized(oauth_token):
    if oauth_token is None:
        flash("Authorization failed.")
        return redirect(url_for('index'))
    g.user = User(oauth_token)
    username = github.get('user')['login']
    g.user.username = username
    user = User.query.filter_by(username=username).first()
    if user is None:
        invite = Invited.query.filter_by(username=username).first()
        if invite is None:
            flash("This service is invite-only")
            return redirect(url_for('manage'))
        user = db_session.merge(g.user)
        g.user = user
        db_session.add(user)
        db_session.commit()
    session['user_id'] = user.user_id
    return redirect(url_for('manage'))

@app.route('/manage', methods=['GET', 'POST'])
def manage():
    user = get_user()
    if user is None:
        flash("Please log in")
        return logout()
    all_repos = github.get('user/repos')
    present = []
    non_present = []
    edit = None
    default_context = 'continuous-integration/travis-ci/push'
    for repo in all_repos:
        on_repo = GithubProjects.query \
            .filter_by(owner=repo['owner']['login'],repo=repo['name']) \
            .first()
        repo_def = {
            "owner": repo['owner']['login'],
            "repo": repo['name'],
            "name": repo['full_name'],
            "id": repo['id'],
        }
        if request.method == 'POST':
            if 'add' in request.form and \
                    int(request.form['add']) == repo['id'] and \
                    on_repo is None:
                return add_repo(repo, default_context)
            elif 'remove' in request.form and \
                    int(request.form['remove']) == repo['id'] and \
                    on_repo is not None:
                return remove_repo(repo, on_repo)
            elif 'edit' in request.form and \
                    int(request.form['edit']) == repo['id'] and \
                    on_repo is not None:
                return edit_repo(on_repo)
        if on_repo is None:
            non_present.append(repo_def)
        else:
            present.append(repo_def)
            if request.args.get('edit') and \
                    request.args.get('edit') == str(repo['id']):
                edit = repo_def
                on_status = GithubStatusPipelines.query \
                    .filter_by(pipeline_id=on_repo.pipeline_id) \
                    .first()
                on_git = GithubGitPipelines.query \
                    .filter_by(pipeline_id=on_repo.pipeline_id) \
                    .first()
                edit['context'] = on_status.context
                edit['master_branch'] = on_git.master_branch
                edit['staging_branch'] = on_git.staging_branch
    return render_template(
        'manage.html',
        username=user.username,
        non_present=non_present,
        present=present,
        edit=edit,
        invite_count=user.invite_count,
        base_url=app.config['BOT_BASEURL'],
    )

def add_repo(repo, context):
    user = get_user()
    # Add repository to our database
    pipeline = Pipeline(None, repo['full_name'])
    db_session.add(pipeline)
    db_session.flush()
    pipeline_id = pipeline.pipeline_id
    on_repo = GithubProjects(
        pipeline_id,
        None,
        repo['owner']['login'],
        repo['name']
    )
    db_session.add(on_repo)
    status = GithubStatusPipelines(
        pipeline_id,
        repo['owner']['login'],
        repo['name'],
        context
    )
    db_session.add(status)
    git = GithubGitPipelines(
        pipeline_id,
        repo['owner']['login'],
        repo['name']
    )
    db_session.add(git)
    db_session.commit()
    # Add our account as a collaborator on Github
    headers = {
        "Content-Type": "application/vnd.github.swamp-thing-preview+json",
    }
    invite=github.request(
        'PUT',
        'repos/' + repo['full_name'] + '/collaborators/' + \
            app.config['BOT_USERNAME'],
        headers=headers,
    )
    headers = {
        "Accept": "token " + app.config['BOT_ACCESS_TOKEN'],
        "Content-Type": "application/vnd.github.swamp-thing-preview+json",
    }
    requests.request(
        'PATCH',
        invite.url,
        headers=headers,
    )
    # Add our webhooks
    github.post(
        'repos/' + repo['full_name'] + '/hooks',
        {
            "name": "web",
            "active": True,
            "config": {
                "url": app.config['BOT_BASEURL'] + '/github-notice',
                "content_type": "json",
                "secret": app.config['BOT_NOTICE_SECRET']
            },
            "events": [ "issue_comment", "pull_request", "team_add" ]
        }
    )
    github.post(
        'repos/' + repo['full_name'] + '/hooks',
        {
            "name": "web",
            "active": True,
            "config": {
                "url": app.config['BOT_BASEURL'] + '/github-status',
                "content_type": "json",
                "secret": app.config['BOT_STATUS_SECRET']
            },
            "events": [ "status" ]
        }
    )
    flash("Added successfully")
    return redirect(url_for('manage'))

def remove_repo(repo, on_repo):
    user = get_user()
    # Remove from our database
    pipeline_id = on_repo.pipeline_id
    db_session.delete(on_repo)
    pipeline = Pipeline.query.filter_by(pipeline_id=pipeline_id).first()
    if not pipeline is None: db_session.delete(pipeline)
    status = GithubStatusPipelines.query \
            .filter_by(pipeline_id=pipeline_id) \
            .first();
    if not status is None: db_session.delete(status)
    git = GithubGitPipelines.query \
            .filter_by(pipeline_id=pipeline_id) \
            .first();
    if not git is None: db_session.delete(git)
    db_session.commit()
    # Remove our collaboratorship
    github.raw_request(
        'DELETE',
        'repos/' + repo['full_name'] + '/collaborators/' + \
            app.config['BOT_USERNAME']
    )
    # Remove our webhooks
    hooks = github.get(
        'repos/' + repo['full_name'] + '/hooks'
    )
    my_urls = {
        app.config['BOT_BASEURL'] + '/github-notice',
        app.config['BOT_BASEURL'] + '/github-status',
    }
    for webhook in hooks:
        if webhook['name'] == 'web' and webhook['config']['url'] in my_urls:
            github.request(
                'DELETE',
                'repos/' + repo['full_name'] + '/hooks/' + str(webhook['id'])
            )
    flash("Deleted successfully")
    return redirect(url_for('manage'))

def edit_repo(project):
    user = get_user()
    # Remove from our database
    pipeline_id = project.pipeline_id
    status = GithubStatusPipelines.query \
            .filter_by(pipeline_id=pipeline_id) \
            .first();
    status.context = request.form['context']
    git = GithubGitPipelines.query \
            .filter_by(pipeline_id=pipeline_id) \
            .first();
    git.master_branch = request.form['master_branch']
    git.staging_branch = request.form['staging_branch']
    db_session.commit()
    flash("Saved successfully")
    return redirect(url_for('manage'))

@app.route('/invite', methods=['POST'])
def invite():
    user = get_user()
    if user is None:
        flash("Please log in")
        return redirect(url_for('index'))
    if user.invite_count <= 0:
        flash("You're out of invites")
        return redirect(url_for('manage'))
    invite = Invited.query.filter_by(username=request.form['username'])
    if invite.first() is not None:
        flash("This person is already invited")
        return redirect(url_for('manage'))
    invite = Invited(request.form['username'])
    db_session.add(invite)
    user.invite_count = user.invite_count - 1
    db_session.commit()
    flash("Invitation recorded; now let them know!")
    return redirect(url_for('manage'))
<|endoftext|>


<commit_before>'''
This module contains all fo the routines needed to set up a master server, this
involves preparing the three listeners and the workers needed by the master.
'''
# Import python modules
import os
import shutil
import hashlib
import logging
import threading
import multiprocessing
import time
import datetime
import cPickle as pickle
# Import zeromq
import zmq
# Import salt modules
import salt.utils
import salt.crypt
import salt.payload
import salt.client
# Import cryptography modules
from M2Crypto import RSA

log = logging.getLogger(__name__)

class Master(object):
    '''
    The salt master server
    '''
    def __init__(self, opts):
        '''
        Create a salt master server instance
        '''
        self.opts = opts

    def _clear_old_jobs(self):
        '''
        Clean out the old jobs
        '''
        if self.opts['keep_jobs'] == 0:
            return
        diff = self.opts['keep_jobs'] * 60 * 60
        keep = int(time.time()) - diff
        jid_root = os.path.join(self.opts['cachedir'], 'jobs')
        for jid in os.listdir(jid_root):
            if int(jid.split('.')[0]) < keep:
                shutil.rmtree(os.path.join(jid_root, jid))

    def start(self):
        '''
        Turn on the master server components
        '''
        log.info('Starting the Salt Master')
        reqserv = ReqServer(self.opts)
        reqserv.run()
        while True:
            self._clear_old_jobs()
            time.sleep(60)

class Publisher(multiprocessing.Process):
    '''
    The publishing interface, a simple zeromq publisher that sends out the
    commands.
    '''
    def __init__(self, opts):
        multiprocessing.Process.__init__(self)
        self.opts = opts

    def run(self):
        '''
        Bind to the interface specified in the configuration file
        '''
        context = zmq.Context(1)
        pub_sock = context.socket(zmq.PUB)
        pull_sock = context.socket(zmq.PULL)
        pub_uri = 'tcp://%(interface)s:%(publish_port)s' % self.opts
        pull_uri = 'tcp://127.0.0.1:%(publish_pull_port)s' % self.opts
        log.info('Starting the Salt Publisher on %s', pub_uri)
        pub_sock.bind(pub_uri)
        pull_sock.bind(pull_uri)

        while True:
            package = pull_sock.recv()
            log.info('Publishing command')
            pub_sock.send(package)

class ReqServer():
    '''
    Starts up the master request server, minions send results to this
    interface.
    '''
    def __init__(self, opts):
        self.opts = opts
        self.master_key = salt.crypt.MasterKeys(self.opts)
        self.context = zmq.Context(1)
        # Prepare the zeromq sockets
        self.uri = 'tcp://%(interface)s:%(ret_port)s' % self.opts
        self.clients = self.context.socket(zmq.XREP)
        self.workers = self.context.socket(zmq.XREQ)
        self.w_uri = 'inproc://workers'
        # Start the publisher
        self.publisher = Publisher(opts)
        self.publisher.start()
        # Prepare the aes key
        self.key = self.__prep_key()
        self.crypticle = salt.crypt.Crypticle(self.opts['aes'])
        # Make a client
        self.local = salt.client.LocalClient(self.opts['conf_file'])

    def __prep_key(self):
        '''
        A key needs to be placed in the filesystem with permissions 0400 so
        clients are required to run as root.
        '''
        log.info('Preparing the root key for local communication')
        keyfile = os.path.join(self.opts['cachedir'], '.root_key')
        key = salt.crypt.Crypticle.generate_key_string()
        if os.path.isfile(keyfile):
            os.chmod(keyfile, 384)
        open(keyfile, 'w+').write(key)
        os.chmod(keyfile, 256)
        return key

    def __worker(self, ind):
        '''
        Starts up a worker thread
        '''
        in_socket = self.context.socket(zmq.REP)
        in_socket.connect(self.w_uri)
        m_worker = MWorker(self.opts,
                ind,
                self.master_key,
                self.key,
                self.crypticle)
        work_port = m_worker.port
        m_worker.start()

        out_socket = self.context.socket(zmq.REQ)
        out_socket.connect('tcp://127.0.0.1:%s' % work_port)

        while True:
            package = in_socket.recv()
            out_socket.send(package)
            ret = out_socket.recv()
            in_socket.send(ret)

    def __bind(self):
        '''
        Binds the reply server
        '''
        log.info('Setting up the master communication server')
        self.clients.bind(self.uri)

        self.workers.bind(self.w_uri)

        for ind in range(int(self.opts['worker_threads'])):
            threading.Thread(target=lambda: self.__worker(ind)).start()
            time.sleep(0.1)

        zmq.device(zmq.QUEUE, self.clients, self.workers)

    def run(self):
        '''
        Start up the ReqServer
        '''
        self.__bind()

class MWorker(multiprocessing.Process):
    '''
    The worker multiprocess instance to manage the backend operations for the
    salt master.
    '''
    def __init__(self, opts, ind, mkey, key, crypticle):
        multiprocessing.Process.__init__(self)
        self.opts = opts
        self.master_key = mkey
        self.key = key
        self.crypticle = crypticle
        self.port = str(ind + int(self.opts['worker_start_port']))

    def __bind(self):
        '''
        Bind to the local port
        '''
        context = zmq.Context(1)
        socket = context.socket(zmq.REP)
        socket.bind('tcp://127.0.0.1:%s' % self.port)

        while True:
            package = socket.recv()
            payload = salt.payload.unpackage(package)
            ret = salt.payload.package(self._handle_payload(payload))
            socket.send(ret)

    def _prep_jid(self, load):
        '''
        Parses the job return directory, generates a job id and sets up the
        job id directory.
        '''
        jid_root = os.path.join(self.opts['cachedir'], 'jobs')
        jid = datetime.datetime.strftime(
            datetime.datetime.now(), '%Y%m%d%H%M%S%f'
        )
        jid_dir = os.path.join(jid_root, jid)
        if not os.path.isdir(jid_dir):
            os.makedirs(jid_dir)
            pickle.dump(load, open(os.path.join(jid_dir, '.load.p'), 'w+'))
        else:
            return self._prep_jid(load)
        return jid

    def _handle_payload(self, payload):
        '''
        The _handle_payload method is the key method used to figure out what
        needs to be done with communication to the server
        '''
        return {'aes': self._handle_aes,
                'pub': self._handle_pub,
                'clear': self._handle_clear}[payload['enc']](payload['load'])

    def _handle_clear(self, load):
        '''
        Take care of a cleartext command
        '''
        log.info('Clear payload received with command %(cmd)s', load)
        return getattr(self, load['cmd'])(load)

    def _handle_pub(self, load):
        '''
        Handle a command sent via a public key pair
        '''
        log.info('Pubkey payload received with command %(cmd)s', load)

    def _handle_aes(self, load):
        '''
        Handle a command sent via an aes key
        '''
        data = self.crypticle.loads(load)
        log.info('AES payload received with command %(cmd)s', load)
        return getattr(self, data['cmd'])(data)

    def _auth(self, load):
        '''
        Authenticate the client, use the sent public key to encrypt the aes key
        which was generated at start up
        '''
        # 1. Verify that the key we are receiving matches the stored key
        # 2. Store the key if it is not there
        # 3. make an rsa key with the pub key
        # 4. encrypt the aes key as an encrypted pickle
        # 5. package the return and return it
        log.info('Authentication request from %(id)s', load)
        pubfn = os.path.join(self.opts['pki_dir'],
                'minions',
                load['id'])
        pubfn_pend = os.path.join(self.opts['pki_dir'],
                'minions_pre',
                load['id'])
        if self.opts['open_mode']:
            # open mode is turned on, nuts to checks and overwrite whatever
            # is there
            pass
        elif os.path.isfile(pubfn):
            # The key has been accepted check it
            if not open(pubfn, 'r').read() == load['pub']:
                log.error(
                    'Authentication attempt from %(id)s failed, the public '
                    'keys did not match. This may be an attempt to compromise '
                    'the Salt cluster.', load
                )
                ret = {'enc': 'clear',
                       'load': {'ret': False}}
                return ret
        elif not os.path.isfile(pubfn_pend)\
                and not self.opts['auto_accept']:
            # This is a new key, stick it in pre
            log.info('New public key placed in pending for %(id)s', load)
            open(pubfn_pend, 'w+').write(load['pub'])
            ret = {'enc': 'clear',
                   'load': {'ret': True}}
            return ret
        elif os.path.isfile(pubfn_pend)\
                and not self.opts['auto_accept']:
            # This key is in pending, if it is the same key ret True, else
            # ret False
            if not open(pubfn_pend, 'r').read() == load['pub']:
                log.error(
                    'Authentication attempt from %(id)s failed, the public '
                    'keys in pending did not match. This may be an attempt to '
                    'compromise the Salt cluster.', load
                )
                return {'enc': 'clear',
                        'load': {'ret': False}}
            else:
                log.info(
                    'Authentication failed from host %(id)s, the key is in '
                    'pending and needs to be accepted with saltkey -a %(id)s',
                    load
                )
                return {'enc': 'clear',
                        'load': {'ret': True}}
        elif not os.path.isfile(pubfn_pend)\
                and self.opts['auto_accept']:
            # This is a new key and auto_accept is turned on
            pass
        else:
            # Something happened that I have not accounted for, FAIL!
            return {'enc': 'clear',
                    'load': {'ret': False}}

        log.info('Authentication accepted from %(id)s', load)
        open(pubfn, 'w+').write(load['pub'])
        key = RSA.load_pub_key(pubfn)
        ret = {'enc': 'pub',
               'pub_key': self.master_key.pub_str,
               'token': self.master_key.token,
               'publish_port': self.opts['publish_port'],
              }
        ret['aes'] = key.public_encrypt(self.opts['aes'], 4)
        if self.opts['cluster_masters']:
            self._send_cluster()
        return ret

    def _serve_file(self, load):
        '''
        Return a chunk from a file based on the data received
        '''
        if not load.has_key('path') or not load.has_key('loc'):
            return False
        path = load['path']
        if path.startswith('/'):
            path = load['path'][1:]
        path = os.path.join(self.opts['file_root'], path)
        if not os.path.isfile(path):
            return ''
        fn_ = open(path, 'rb')
        fn_.seek(load['loc'])
        return self.crypticle.dumps(fn_.read(self.opts['file_buffer_size']))

    def _file_hash(self, load):
        '''
        Return a file hash, the hash type is set in the master config file
        '''
        if not load.has_key('path'):
            return False
        path = os.path.join(self.opts['file_root'], load['path'])
        if not os.path.isfile(path):
            return False
        ret = {}
        ret['hsum'] = getattr(hashlib, self.opts['hash_type'])(open(path,
            'rb').read()).hexdigest()
        ret['hash_type'] = self.opts['hash_type']
        return self.crypticle.dumps(ret)

    def _return(self, load):
        '''
        Handle the return data sent from the minions
        '''
        # If the return data is invalid, just ignore it
        if not load.has_key('return')\
                or not load.has_key('jid')\
                or not load.has_key('id'):
            return False
        log.info('Got return from %(id)s for job %(jid)s', load)
        jid_dir = os.path.join(self.opts['cachedir'], 'jobs', load['jid'])
        if not os.path.isdir(jid_dir):
            log.error(
                'An inconsistency occurred, a job was received with a job id '
                'that is not present on the master: %(jid)s', load
            )
            return False
        hn_dir = os.path.join(jid_dir, load['id'])
        if not os.path.isdir(hn_dir):
            os.makedirs(hn_dir)
        pickle.dump(load['return'],
                open(os.path.join(hn_dir, 'return.p'), 'w+'))

    def _send_cluster(self):
        '''
        Send the cluster data out
        '''
        log.debug('Sending out cluster data')
        ret = self.local.cmd(self.opts['cluster_masters'],
                'cluster.distrib',
                self._cluster_load(),
                0,
                'list'
                )
        log.debug('Cluster distributed: %s', ret)

    def _cluster_load(self):
        '''
        Generates the data sent to the cluster nodes.
        '''
        minions = {}
        master_pem = ''
        master_conf = open(self.opts['conf_file'], 'r').read()
        minion_dir = os.path.join(self.opts['pki_dir'], 'minions')
        for host in os.listdir(minion_dir):
            pub = os.path.join(minion_dir, host)
            minions[host] = open(pub, 'r').read()
        if self.opts['cluster_mode'] == 'full':
            master_pem = open(os.path.join(self.opts['pki_dir'],
                'master.pem')).read()
        return [minions,
                master_conf,
                master_pem,
                self.opts['conf_file']]

    def publish(self, clear_load):
        '''
        This method sends out publications to the minions
        '''
        if not clear_load.pop('key') == self.key:
            return ''
        jid = self._prep_jid(clear_load)
        payload = {'enc': 'aes'}
        load = {
                'fun': clear_load['fun'],
                'arg': clear_load['arg'],
                'tgt': clear_load['tgt'],
                'jid': jid,
                'ret': clear_load['ret'],
               }
        if clear_load.has_key('tgt_type'):
            load['tgt_type'] = clear_load['tgt_type']
        payload['load'] = self.crypticle.dumps(load)
        context = zmq.Context(1)
        pub_sock = context.socket(zmq.PUSH)
        pub_sock.connect('tcp://127.0.0.1:%(publish_pull_port)s' % self.opts)
        pub_sock.send(salt.payload.package(payload))
        return {'enc': 'clear',
                'load': {'jid': jid}}

    def run(self):
        '''
        Start a Master Worker
        '''
        self.__bind()
<commit_msg>Allow the minion to retrive the master's configuration options<commit_after>'''
This module contains all fo the routines needed to set up a master server, this
involves preparing the three listeners and the workers needed by the master.
'''
# Import python modules
import os
import shutil
import hashlib
import logging
import threading
import multiprocessing
import time
import datetime
import cPickle as pickle
# Import zeromq
import zmq
# Import salt modules
import salt.utils
import salt.crypt
import salt.payload
import salt.client
# Import cryptography modules
from M2Crypto import RSA

log = logging.getLogger(__name__)

class Master(object):
    '''
    The salt master server
    '''
    def __init__(self, opts):
        '''
        Create a salt master server instance
        '''
        self.opts = opts

    def _clear_old_jobs(self):
        '''
        Clean out the old jobs
        '''
        if self.opts['keep_jobs'] == 0:
            return
        diff = self.opts['keep_jobs'] * 60 * 60
        keep = int(time.time()) - diff
        jid_root = os.path.join(self.opts['cachedir'], 'jobs')
        for jid in os.listdir(jid_root):
            if int(jid.split('.')[0]) < keep:
                shutil.rmtree(os.path.join(jid_root, jid))

    def start(self):
        '''
        Turn on the master server components
        '''
        log.info('Starting the Salt Master')
        reqserv = ReqServer(self.opts)
        reqserv.run()
        while True:
            self._clear_old_jobs()
            time.sleep(60)

class Publisher(multiprocessing.Process):
    '''
    The publishing interface, a simple zeromq publisher that sends out the
    commands.
    '''
    def __init__(self, opts):
        multiprocessing.Process.__init__(self)
        self.opts = opts

    def run(self):
        '''
        Bind to the interface specified in the configuration file
        '''
        context = zmq.Context(1)
        pub_sock = context.socket(zmq.PUB)
        pull_sock = context.socket(zmq.PULL)
        pub_uri = 'tcp://%(interface)s:%(publish_port)s' % self.opts
        pull_uri = 'tcp://127.0.0.1:%(publish_pull_port)s' % self.opts
        log.info('Starting the Salt Publisher on %s', pub_uri)
        pub_sock.bind(pub_uri)
        pull_sock.bind(pull_uri)

        while True:
            package = pull_sock.recv()
            log.info('Publishing command')
            pub_sock.send(package)

class ReqServer():
    '''
    Starts up the master request server, minions send results to this
    interface.
    '''
    def __init__(self, opts):
        self.opts = opts
        self.master_key = salt.crypt.MasterKeys(self.opts)
        self.context = zmq.Context(1)
        # Prepare the zeromq sockets
        self.uri = 'tcp://%(interface)s:%(ret_port)s' % self.opts
        self.clients = self.context.socket(zmq.XREP)
        self.workers = self.context.socket(zmq.XREQ)
        self.w_uri = 'inproc://workers'
        # Start the publisher
        self.publisher = Publisher(opts)
        self.publisher.start()
        # Prepare the aes key
        self.key = self.__prep_key()
        self.crypticle = salt.crypt.Crypticle(self.opts['aes'])
        # Make a client
        self.local = salt.client.LocalClient(self.opts['conf_file'])

    def __prep_key(self):
        '''
        A key needs to be placed in the filesystem with permissions 0400 so
        clients are required to run as root.
        '''
        log.info('Preparing the root key for local communication')
        keyfile = os.path.join(self.opts['cachedir'], '.root_key')
        key = salt.crypt.Crypticle.generate_key_string()
        if os.path.isfile(keyfile):
            os.chmod(keyfile, 384)
        open(keyfile, 'w+').write(key)
        os.chmod(keyfile, 256)
        return key

    def __worker(self, ind):
        '''
        Starts up a worker thread
        '''
        in_socket = self.context.socket(zmq.REP)
        in_socket.connect(self.w_uri)
        m_worker = MWorker(self.opts,
                ind,
                self.master_key,
                self.key,
                self.crypticle)
        work_port = m_worker.port
        m_worker.start()

        out_socket = self.context.socket(zmq.REQ)
        out_socket.connect('tcp://127.0.0.1:%s' % work_port)

        while True:
            package = in_socket.recv()
            out_socket.send(package)
            ret = out_socket.recv()
            in_socket.send(ret)

    def __bind(self):
        '''
        Binds the reply server
        '''
        log.info('Setting up the master communication server')
        self.clients.bind(self.uri)

        self.workers.bind(self.w_uri)

        for ind in range(int(self.opts['worker_threads'])):
            threading.Thread(target=lambda: self.__worker(ind)).start()
            time.sleep(0.1)

        zmq.device(zmq.QUEUE, self.clients, self.workers)

    def run(self):
        '''
        Start up the ReqServer
        '''
        self.__bind()

class MWorker(multiprocessing.Process):
    '''
    The worker multiprocess instance to manage the backend operations for the
    salt master.
    '''
    def __init__(self, opts, ind, mkey, key, crypticle):
        multiprocessing.Process.__init__(self)
        self.opts = opts
        self.master_key = mkey
        self.key = key
        self.crypticle = crypticle
        self.port = str(ind + int(self.opts['worker_start_port']))

    def __bind(self):
        '''
        Bind to the local port
        '''
        context = zmq.Context(1)
        socket = context.socket(zmq.REP)
        socket.bind('tcp://127.0.0.1:%s' % self.port)

        while True:
            package = socket.recv()
            payload = salt.payload.unpackage(package)
            ret = salt.payload.package(self._handle_payload(payload))
            socket.send(ret)

    def _prep_jid(self, load):
        '''
        Parses the job return directory, generates a job id and sets up the
        job id directory.
        '''
        jid_root = os.path.join(self.opts['cachedir'], 'jobs')
        jid = datetime.datetime.strftime(
            datetime.datetime.now(), '%Y%m%d%H%M%S%f'
        )
        jid_dir = os.path.join(jid_root, jid)
        if not os.path.isdir(jid_dir):
            os.makedirs(jid_dir)
            pickle.dump(load, open(os.path.join(jid_dir, '.load.p'), 'w+'))
        else:
            return self._prep_jid(load)
        return jid

    def _handle_payload(self, payload):
        '''
        The _handle_payload method is the key method used to figure out what
        needs to be done with communication to the server
        '''
        return {'aes': self._handle_aes,
                'pub': self._handle_pub,
                'clear': self._handle_clear}[payload['enc']](payload['load'])

    def _handle_clear(self, load):
        '''
        Take care of a cleartext command
        '''
        log.info('Clear payload received with command %(cmd)s', load)
        return getattr(self, load['cmd'])(load)

    def _handle_pub(self, load):
        '''
        Handle a command sent via a public key pair
        '''
        log.info('Pubkey payload received with command %(cmd)s', load)

    def _handle_aes(self, load):
        '''
        Handle a command sent via an aes key
        '''
        data = self.crypticle.loads(load)
        log.info('AES payload received with command %(cmd)s', load)
        return getattr(self, data['cmd'])(data)

    def _auth(self, load):
        '''
        Authenticate the client, use the sent public key to encrypt the aes key
        which was generated at start up
        '''
        # 1. Verify that the key we are receiving matches the stored key
        # 2. Store the key if it is not there
        # 3. make an rsa key with the pub key
        # 4. encrypt the aes key as an encrypted pickle
        # 5. package the return and return it
        log.info('Authentication request from %(id)s', load)
        pubfn = os.path.join(self.opts['pki_dir'],
                'minions',
                load['id'])
        pubfn_pend = os.path.join(self.opts['pki_dir'],
                'minions_pre',
                load['id'])
        if self.opts['open_mode']:
            # open mode is turned on, nuts to checks and overwrite whatever
            # is there
            pass
        elif os.path.isfile(pubfn):
            # The key has been accepted check it
            if not open(pubfn, 'r').read() == load['pub']:
                log.error(
                    'Authentication attempt from %(id)s failed, the public '
                    'keys did not match. This may be an attempt to compromise '
                    'the Salt cluster.', load
                )
                ret = {'enc': 'clear',
                       'load': {'ret': False}}
                return ret
        elif not os.path.isfile(pubfn_pend)\
                and not self.opts['auto_accept']:
            # This is a new key, stick it in pre
            log.info('New public key placed in pending for %(id)s', load)
            open(pubfn_pend, 'w+').write(load['pub'])
            ret = {'enc': 'clear',
                   'load': {'ret': True}}
            return ret
        elif os.path.isfile(pubfn_pend)\
                and not self.opts['auto_accept']:
            # This key is in pending, if it is the same key ret True, else
            # ret False
            if not open(pubfn_pend, 'r').read() == load['pub']:
                log.error(
                    'Authentication attempt from %(id)s failed, the public '
                    'keys in pending did not match. This may be an attempt to '
                    'compromise the Salt cluster.', load
                )
                return {'enc': 'clear',
                        'load': {'ret': False}}
            else:
                log.info(
                    'Authentication failed from host %(id)s, the key is in '
                    'pending and needs to be accepted with saltkey -a %(id)s',
                    load
                )
                return {'enc': 'clear',
                        'load': {'ret': True}}
        elif not os.path.isfile(pubfn_pend)\
                and self.opts['auto_accept']:
            # This is a new key and auto_accept is turned on
            pass
        else:
            # Something happened that I have not accounted for, FAIL!
            return {'enc': 'clear',
                    'load': {'ret': False}}

        log.info('Authentication accepted from %(id)s', load)
        open(pubfn, 'w+').write(load['pub'])
        key = RSA.load_pub_key(pubfn)
        ret = {'enc': 'pub',
               'pub_key': self.master_key.pub_str,
               'token': self.master_key.token,
               'publish_port': self.opts['publish_port'],
              }
        ret['aes'] = key.public_encrypt(self.opts['aes'], 4)
        if self.opts['cluster_masters']:
            self._send_cluster()
        return ret

    def _serve_file(self, load):
        '''
        Return a chunk from a file based on the data received
        '''
        if not load.has_key('path') or not load.has_key('loc'):
            return False
        path = load['path']
        if path.startswith('/'):
            path = load['path'][1:]
        path = os.path.join(self.opts['file_root'], path)
        if not os.path.isfile(path):
            return ''
        fn_ = open(path, 'rb')
        fn_.seek(load['loc'])
        return self.crypticle.dumps(fn_.read(self.opts['file_buffer_size']))

    def _file_hash(self, load):
        '''
        Return a file hash, the hash type is set in the master config file
        '''
        if not load.has_key('path'):
            return False
        path = os.path.join(self.opts['file_root'], load['path'])
        if not os.path.isfile(path):
            return False
        ret = {}
        ret['hsum'] = getattr(hashlib, self.opts['hash_type'])(open(path,
            'rb').read()).hexdigest()
        ret['hash_type'] = self.opts['hash_type']
        return self.crypticle.dumps(ret)

    def _master_opts(self, load):
        '''
        Return the master options to the minion
        '''
        return self.crypticle.dumps(self.opts)

    def _return(self, load):
        '''
        Handle the return data sent from the minions
        '''
        # If the return data is invalid, just ignore it
        if not load.has_key('return')\
                or not load.has_key('jid')\
                or not load.has_key('id'):
            return False
        log.info('Got return from %(id)s for job %(jid)s', load)
        jid_dir = os.path.join(self.opts['cachedir'], 'jobs', load['jid'])
        if not os.path.isdir(jid_dir):
            log.error(
                'An inconsistency occurred, a job was received with a job id '
                'that is not present on the master: %(jid)s', load
            )
            return False
        hn_dir = os.path.join(jid_dir, load['id'])
        if not os.path.isdir(hn_dir):
            os.makedirs(hn_dir)
        pickle.dump(load['return'],
                open(os.path.join(hn_dir, 'return.p'), 'w+'))

    def _send_cluster(self):
        '''
        Send the cluster data out
        '''
        log.debug('Sending out cluster data')
        ret = self.local.cmd(self.opts['cluster_masters'],
                'cluster.distrib',
                self._cluster_load(),
                0,
                'list'
                )
        log.debug('Cluster distributed: %s', ret)

    def _cluster_load(self):
        '''
        Generates the data sent to the cluster nodes.
        '''
        minions = {}
        master_pem = ''
        master_conf = open(self.opts['conf_file'], 'r').read()
        minion_dir = os.path.join(self.opts['pki_dir'], 'minions')
        for host in os.listdir(minion_dir):
            pub = os.path.join(minion_dir, host)
            minions[host] = open(pub, 'r').read()
        if self.opts['cluster_mode'] == 'full':
            master_pem = open(os.path.join(self.opts['pki_dir'],
                'master.pem')).read()
        return [minions,
                master_conf,
                master_pem,
                self.opts['conf_file']]

    def publish(self, clear_load):
        '''
        This method sends out publications to the minions
        '''
        if not clear_load.pop('key') == self.key:
            return ''
        jid = self._prep_jid(clear_load)
        payload = {'enc': 'aes'}
        load = {
                'fun': clear_load['fun'],
                'arg': clear_load['arg'],
                'tgt': clear_load['tgt'],
                'jid': jid,
                'ret': clear_load['ret'],
               }
        if clear_load.has_key('tgt_type'):
            load['tgt_type'] = clear_load['tgt_type']
        payload['load'] = self.crypticle.dumps(load)
        context = zmq.Context(1)
        pub_sock = context.socket(zmq.PUSH)
        pub_sock.connect('tcp://127.0.0.1:%(publish_pull_port)s' % self.opts)
        pub_sock.send(salt.payload.package(payload))
        return {'enc': 'clear',
                'load': {'jid': jid}}

    def run(self):
        '''
        Start a Master Worker
        '''
        self.__bind()
<|endoftext|>


<commit_before>#!/usr/bin/python3
# This file is part of Cockpit.
#
# Copyright (C) 2021 Red Hat, Inc.
#
# Cockpit is free software; you can redistribute it and/or modify it
# under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation; either version 2.1 of the License, or
# (at your option) any later version.
#
# Cockpit is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Cockpit; If not, see <http://www.gnu.org/licenses/>.

import glob
import io
import multiprocessing
import os
import urllib
import subprocess
import sys
import zipfile

def message(*args):
    print(*args, file=sys.stderr)

def build_dist():
    '''Build a dist tarball for CI testing

    This supports completely clean git trees, unpacked release tarballs, and already configured trees.
    Returns path to built tarball.
    '''
    if not os.path.exists("Makefile"):
        if os.path.exists('./configure'):
            # unconfigured release tarball
            subprocess.check_call('./configure')
        else:
            # clean git checkout
            subprocess.check_call('./autogen.sh')

    # this is for a development build, not a release, so we care about speed, not best size
    subprocess.check_call(["make", "--silent", "-j%i" % multiprocessing.cpu_count(),
                           "NO_DIST_CACHE=1", "XZ_COMPRESS_FLAGS=-0", "dist"])
    return subprocess.check_output(["make", "dump-dist"], universal_newlines=True).strip()

def download_dist():
    '''Download dists tarball for current git SHA from GitHub

    These are produced by .github/workflows/build-dist.yml for every PR and push.
    This is a lot faster than having to npm install and run webpack.

    Returns path to downloaded tarball, or None if it isn't available.
    This can happen because the current directory is not a git checkout, or it is
    a SHA which is not pushed/PRed, or there is no ~/.config/github-token available.
    '''
    try:
        sha = subprocess.check_output(["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL).decode().strip()
    except subprocess.CalledProcessError:
        message("make_dist: not a git repository")
        return None

    if subprocess.call(["git", "diff", "--quiet", "--", ":^test"]) > 0:
        message("make_dist: uncommitted local changes, skipping download")
        return None

    dists = glob.glob(f"cockpit-*{sha[:8]}*.tar.xz")
    if dists:
        message("make_dist: already downloaded", dists[0])
        return os.path.abspath(dists[0])

    # autogen.sh does that for build_dist()
    if not os.path.exists("bots"):
        subprocess.check_call(['tools/make-bots'])

    sys.path.append(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "bots"))
    import task

    # task.api defaults to the current checkout's origin, but the artifacts are on cockpit-project
    # except if a developer wants to change the artifact building workflow on their fork, support that with $GITHUB_BASE
    api = task.github.GitHub(repo=os.getenv("GITHUB_BASE", "cockpit-project/cockpit"))

    # downloading GitHub artifacts requires a token
    if not api.token:
        message("make_dist: no GitHub API token available")
        return None

    # iterate artifacts and search our SHA
    page = 1
    batch_size = 100
    download_url = None
    while not download_url and batch_size == 100:
        batch = api.get(f"actions/artifacts?per_page={batch_size}&page={page}")["artifacts"]
        for artifact in batch:
            if artifact["name"] == "dist-" + sha:
                download_url = artifact["archive_download_url"]
                break
        # if the current batch is < 100, we have the last page and can stop
        batch_size = len(batch)
        page += 1

    if not download_url:
        message(f"make_dist: no download available for commit {sha}")
        return None

    sys.stderr.write(f"make_dist: Downloading dist tarball from {download_url} ...\n")
    request = urllib.request.Request(download_url, headers={"Authorization": "token " + api.token})
    zipio = io.BytesIO()
    try:
        with urllib.request.urlopen(request) as response:
            if os.isatty(sys.stderr.fileno()):
                total_size = 0
            else:
                total_size = None
            MB = 10**6
            # read zip into a stringio, as the stream is not seekable and zip requires that
            while True:
                block = response.read(MB)
                if len(block) == 0:
                    break
                if total_size is not None:
                    total_size += len(block)
                    sys.stderr.write(f"\r{ total_size // MB } MB")

                zipio.write(block)

            # clear the download progress in tty mode
            if total_size is not None:
                sys.stderr.write("\r                             \r")

    except urllib.error.HTTPError as e:
        message("make_dist: Download failed:", e)
        return None

    with zipfile.ZipFile(zipio) as fzip:
        names = fzip.namelist()
        if len(names) != 1 or not names[0].endswith(".tar.xz"):
            message("make_dist: expected zip artifact with exactly one tar.xz member")
            return None
        tar_path = fzip.extract(names[0])

    # Extract node_modules and dist locally for speeding up the build and allowing integration tests to run
    unpack_dirs = [d for d in ["dist", "node_modules"] if not os.path.exists(d)]
    if unpack_dirs:
        message("make_dist: Extracting directories from tarball:", ' '.join(unpack_dirs))
        prefix = os.path.basename(tar_path).split('.tar')[0] + '/'
        prefixed_unpack_dirs = [prefix + d for d in unpack_dirs]
        subprocess.check_call(["tar", "--touch", "--strip-components=1", "-xf", tar_path] + prefixed_unpack_dirs)

    return tar_path

def make_dist():
    # first try to download a pre-generated dist tarball; this is a lot faster
    # but these tarballs are built for production NPM mode
    source = None
    if os.getenv("NODE_ENV") != "development":
        source = download_dist()
    if not source:
        source = build_dist()
    return source

if __name__ == '__main__':
    print(make_dist())
<commit_msg>test: extend make_dist.py to have download-only option<commit_after>#!/usr/bin/python3
# This file is part of Cockpit.
#
# Copyright (C) 2021 Red Hat, Inc.
#
# Cockpit is free software; you can redistribute it and/or modify it
# under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation; either version 2.1 of the License, or
# (at your option) any later version.
#
# Cockpit is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Cockpit; If not, see <http://www.gnu.org/licenses/>.

import glob
import io
import multiprocessing
import os
import urllib
import subprocess
import sys
import zipfile
import argparse

def message(*args):
    print(*args, file=sys.stderr)

def build_dist():
    '''Build a dist tarball for CI testing

    This supports completely clean git trees, unpacked release tarballs, and already configured trees.
    Returns path to built tarball.
    '''
    if not os.path.exists("Makefile"):
        if os.path.exists('./configure'):
            # unconfigured release tarball
            subprocess.check_call('./configure')
        else:
            # clean git checkout
            subprocess.check_call('./autogen.sh')

    # this is for a development build, not a release, so we care about speed, not best size
    subprocess.check_call(["make", "--silent", "-j%i" % multiprocessing.cpu_count(),
                           "NO_DIST_CACHE=1", "XZ_COMPRESS_FLAGS=-0", "dist"])
    return subprocess.check_output(["make", "dump-dist"], universal_newlines=True).strip()

def download_dist():
    '''Download dists tarball for current git SHA from GitHub

    These are produced by .github/workflows/build-dist.yml for every PR and push.
    This is a lot faster than having to npm install and run webpack.

    Returns path to downloaded tarball, or None if it isn't available.
    This can happen because the current directory is not a git checkout, or it is
    a SHA which is not pushed/PRed, or there is no ~/.config/github-token available.
    '''
    try:
        sha = subprocess.check_output(["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL).decode().strip()
    except subprocess.CalledProcessError:
        message("make_dist: not a git repository")
        return None

    if subprocess.call(["git", "diff", "--quiet", "--", ":^test"]) > 0:
        message("make_dist: uncommitted local changes, skipping download")
        return None

    dists = glob.glob(f"cockpit-*{sha[:8]}*.tar.xz")
    if dists:
        message("make_dist: already downloaded", dists[0])
        return os.path.abspath(dists[0])

    # autogen.sh does that for build_dist()
    if not os.path.exists("bots"):
        subprocess.check_call(['tools/make-bots'])

    sys.path.append(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "bots"))
    import task

    # task.api defaults to the current checkout's origin, but the artifacts are on cockpit-project
    # except if a developer wants to change the artifact building workflow on their fork, support that with $GITHUB_BASE
    api = task.github.GitHub(repo=os.getenv("GITHUB_BASE", "cockpit-project/cockpit"))

    # downloading GitHub artifacts requires a token
    if not api.token:
        message("make_dist: no GitHub API token available")
        return None

    # iterate artifacts and search our SHA
    page = 1
    batch_size = 100
    download_url = None
    while not download_url and batch_size == 100:
        batch = api.get(f"actions/artifacts?per_page={batch_size}&page={page}")["artifacts"]
        for artifact in batch:
            if artifact["name"] == "dist-" + sha:
                download_url = artifact["archive_download_url"]
                break
        # if the current batch is < 100, we have the last page and can stop
        batch_size = len(batch)
        page += 1

    if not download_url:
        message(f"make_dist: no download available for commit {sha}")
        return None

    sys.stderr.write(f"make_dist: Downloading dist tarball from {download_url} ...\n")
    request = urllib.request.Request(download_url, headers={"Authorization": "token " + api.token})
    zipio = io.BytesIO()
    try:
        with urllib.request.urlopen(request) as response:
            if os.isatty(sys.stderr.fileno()):
                total_size = 0
            else:
                total_size = None
            MB = 10**6
            # read zip into a stringio, as the stream is not seekable and zip requires that
            while True:
                block = response.read(MB)
                if len(block) == 0:
                    break
                if total_size is not None:
                    total_size += len(block)
                    sys.stderr.write(f"\r{ total_size // MB } MB")

                zipio.write(block)

            # clear the download progress in tty mode
            if total_size is not None:
                sys.stderr.write("\r                             \r")

    except urllib.error.HTTPError as e:
        message("make_dist: Download failed:", e)
        return None

    with zipfile.ZipFile(zipio) as fzip:
        names = fzip.namelist()
        if len(names) != 1 or not names[0].endswith(".tar.xz"):
            message("make_dist: expected zip artifact with exactly one tar.xz member")
            return None
        tar_path = fzip.extract(names[0])

    # Extract node_modules and dist locally for speeding up the build and allowing integration tests to run
    unpack_dirs = [d for d in ["dist", "node_modules"] if not os.path.exists(d)]
    if unpack_dirs:
        message("make_dist: Extracting directories from tarball:", ' '.join(unpack_dirs))
        prefix = os.path.basename(tar_path).split('.tar')[0] + '/'
        prefixed_unpack_dirs = [prefix + d for d in unpack_dirs]
        subprocess.check_call(["tar", "--touch", "--strip-components=1", "-xf", tar_path] + prefixed_unpack_dirs)

    return tar_path

def make_dist(download_only=False):
    # first try to download a pre-generated dist tarball; this is a lot faster
    # but these tarballs are built for production NPM mode
    source = None
    if os.getenv("NODE_ENV") != "development":
        source = download_dist()
    if not source:
        if not download_only:
            source = build_dist()
        else:
            print("make_dist: Download failed: artifact does not exist")
            sys.exit(1)
    return source

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Download or build release tarbal")
    parser.add_argument('-d', '--download-only', action='store_true', help="Fail instead of build locally if download is not available")
    args = parser.parse_args()
    print(make_dist(args.download_only))
<|endoftext|>


<commit_before>'''
Routines to set up a minion
'''
# Import python libs
import os
import distutils.sysconfig
import glob
import re
import time
import logging
import tempfile
import traceback
import shutil
import threading
import multiprocessing

# Import zeromq libs
import zmq
# Import salt libs
import salt.crypt
from salt.crypt import AuthenticationError
import salt.utils
import salt.modules
import salt.returners
import salt.loader

log = logging.getLogger(__name__)

# To set up a minion:
# 1, Read in the configuration
# 2. Generate the function mapping dict
# 3. Authenticate with the master
# 4. Store the aes key
# 5. connect to the publisher
# 6. handle publications

class MinionError(Exception): pass

class Minion(object):
    '''
    This class instantiates a minion, runs connections for a minion, and loads
    all of the functions into the minion
    '''
    def __init__(self, opts):
        '''
        Pass in the options dict
        '''
        self.opts = opts
        self.mod_opts = self.__prep_mod_opts()
        self.functions, self.returners = self.__load_modules()
        self.matcher = Matcher(self.opts, self.functions)
        self.authenticate()

    def __prep_mod_opts(self):
        '''
        Returns a deep copy of the opts with key bits stripped out
        '''
        mod_opts = {}
        for key, val in self.opts.items():
            if key == 'logger':
                continue
            mod_opts[key] = val
        return mod_opts

    def __load_modules(self):
        '''
        Return the functions and the returners loaded up from the loader module
        '''
        functions = salt.loader.minion_mods(self.opts)
        returners = salt.loader.returners(self.opts)
        return functions, returners

    def _handle_payload(self, payload):
        '''
        Takes a payload from the master publisher and does whatever the
        master wants done.
        '''
        {'aes': self._handle_aes,
         'pub': self._handle_pub,
         'clear': self._handle_clear}[payload['enc']](payload['load'])

    def _handle_aes(self, load):
        '''
        Takes the aes encrypted load, decrypts is and runs the encapsulated
        instructions
        '''
        data = None
        try:
            data = self.crypticle.loads(load)
        except AuthenticationError:
            self.authenticate()
            data = self.crypticle.loads(load)
        # Verify that the publication is valid
        if not data.has_key('tgt')\
                or not data.has_key('jid')\
                or not data.has_key('fun')\
                or not data.has_key('arg'):
            return
        # Verify that the publication applies to this minion
        if data.has_key('tgt_type'):
            if not getattr(self.matcher, data['tgt_type'] + '_match')(data['tgt']):
                return
        else:
            if not self.matcher.glob_match(data['tgt']):
                return
        self._handle_decoded_payload(data)

    def _handle_pub(self, load):
        '''
        Handle public key payloads
        '''
        pass

    def _handle_clear(self, load):
        '''
        Handle un-encrypted transmissions
        '''
        pass

    def _handle_decoded_payload(self, data):
        '''
        Override this method if you wish to handle the decoded data differently.
        '''
        if self.opts['multiprocessing']:
            if type(data['fun']) == type(list()):
                multiprocessing.Process(
                    target=lambda: self._thread_multi_return(data)
                ).start()
            else:
                multiprocessing.Process(
                    target=lambda: self._thread_return(data)
                ).start()
        else:
            if type(data['fun']) == type(list()):
                threading.Thread(
                    target=lambda: self._thread_multi_return(data)
                ).start()
            else:
                threading.Thread(
                    target=lambda: self._thread_return(data)
                ).start()

    def _thread_return(self, data):
        '''
        This method should be used as a threading target, start the actual
        minion side execution.
        '''
        ret = {}
        for ind in range(0, len(data['arg'])):
            try:
                data['arg'][ind] = eval(data['arg'][ind])
            except:
                pass

        try:
            ret['return'] = self.functions[data['fun']](*data['arg'])
        except Exception as exc:
            trb = traceback.format_exc()
            log.warning('The minion function caused an exception: %s', exc)
            ret['return'] = trb
        ret['jid'] = data['jid']
        if data['ret']:
            ret['id'] = self.opts['id']
            try:
                self.returners[data['ret']](ret)
            except Exception as exc:
                log.error('The return failed for job %s %s', data['jid'], exc)
        else:
            self._return_pub(ret)

    def _thread_multi_return(self, data):
        '''
        This method should be used as a threading target, start the actual
        minion side execution.
        '''
        ret = {'return': {}}
        for ind in range(0, len(data['fun'])):
            for index in range(0, len(data['arg'][ind])):
                try:
                    data['arg'][ind][index] = eval(data['arg'][ind][index])
                except:
                    pass

            try:
                ret['return'][data['fun'][ind]]\
                    = self.functions[data['fun'][ind]](*data['arg'][ind])
            except Exception as exc:
                trb = traceback.format_exc()
                log.warning('The minion function caused an exception: %s', exc)
                ret['return'][data['fun'][ind]] = trb
            ret['jid'] = data['jid']
        if data['ret']:
            ret['id'] = self.opts['id']
            try:
                self.returners[data['ret']](ret)
            except Exception as exc:
                log.error('The return failed for job %s %s', data['jid'], exc)
        else:
            self._return_pub(ret)

    def _return_pub(self, ret):
        '''
        Return the data from the executed command to the master server
        '''
        log.info('Returning information for job: %(jid)s', ret)
        context = zmq.Context()
        socket = context.socket(zmq.REQ)
        socket.connect(self.opts['master_uri'])
        payload = {'enc': 'aes'}
        load = {'return': ret['return'],
                'cmd': '_return',
                'jid': ret['jid'],
                'id': self.opts['id']}
        payload['load'] = self.crypticle.dumps(load)
        socket.send_pyobj(payload)
        return socket.recv()

    def reload_functions(self):
        '''
        Reload the functions dict for this minion, reading in any new functions
        '''
        self.functions = self.__load_functions()
        log.debug('Refreshed functions, loaded functions: %s', self.functions)
        return True

    def authenticate(self):
        '''
        Authenticate with the master, this method breaks the functional
        paradigm, it will update the master information from a fresh sign in,
        signing in can occur as often as needed to keep up with the revolving
        master aes key.
        '''
        log.debug('Attempting to authenticate with the Salt Master')
        auth = salt.crypt.Auth(self.opts)
        while True:
            creds = auth.sign_in()
            if creds != 'retry':
                log.info('Authentication with master successful!')
                break
            log.info('Waiting for minion key to be accepted by the master.')
            time.sleep(10)
        self.aes = creds['aes']
        self.publish_port = creds['publish_port']
        self.crypticle = salt.crypt.Crypticle(self.aes)

    def tune_in(self):
        '''
        Lock onto the publisher. This is the main event loop for the minion
        '''
        master_pub = 'tcp://' + self.opts['master'] + ':'\
                   + str(self.publish_port)
        context = zmq.Context()
        socket = context.socket(zmq.SUB)
        socket.connect(master_pub)
        socket.setsockopt(zmq.SUBSCRIBE, '')
        while True:
            payload = socket.recv_pyobj()
            self._handle_payload(payload)

class Matcher(object):
    '''
    Use to return the value for matching calls from the master
    '''
    def __init__(self, opts, functions=None):
        self.opts = opts
        if not functions:
            functions = salt.loader.minion_mods(self.opts)
        else:
            self.functions = functions

    def glob_match(self, tgt):
        '''
        Returns true if the passed glob matches the id
        '''
        tmp_dir = tempfile.mkdtemp()
        cwd = os.getcwd()
        os.chdir(tmp_dir)
        open(self.opts['id'], 'w+').write('salt')
        ret = bool(glob.glob(tgt))
        os.chdir(cwd)
        shutil.rmtree(tmp_dir)
        return ret

    def pcre_match(self, tgt):
        '''
        Returns true if the passed pcre regex matches
        '''
        return bool(re.match(tgt, self.opts['id']))

    def list_match(self, tgt):
        '''
        Determines if this host is on the list
        '''
        return bool(tgt.count(self.opts['id']))

    def grain_match(self, tgt):
        '''
        Reads in the grains regular expression match
        '''
        comps = tgt.split(':')
        return bool(re.match(comps[1], self.opts['grains'][comps[0]]))

    def exsel_match(self, tgt):
        '''
        Runs a function and return the exit code
        '''
        if not self.functions.has_key(tgt):
            return False
        return(self.functions[tgt]())

class FileClient(object):
    '''
    Interact with the salt master file server.
    '''
    def __init__(self, opts):
        self.opts = opts
        self.auth = salt.crypt.SAuth(opts)
        self.socket = self.__get_socket()

    def __get_socket(self):
        '''
        Return the ZeroMQ socket to use
        '''
        context = zmq.Context()
        socket = context.socket(zmq.REQ)
        socket.connect(self.opts['master_uri'])
        return socket

    def _check_proto(self, path):
        '''
        Make sure that this path is intended for the salt master and trim it
        '''
        if not path.startswith('salt://'):
            raise MinionError('Unsupported path')
        return path[:7]

    def get_file(self, path, dest, makedirs=False, env='base'):
        '''
        Get a single file from the salt-master
        '''
        path = self._check_proto(path)
        payload = {'enc': 'aes'}
        destdir = os.path.dirname(dest)
        if not os.path.isdir(destdir):
            if makedirs:
                os.makedirs(destdir)
            else:
                return False
        fn_ = open(dest, 'w+')
        load = {'path': path,
                'env': env,
                'cmd': '_serve_file'}
        while True:
            load['loc'] = fn_.tell()
            payload['load'] = self.crypticle.dumps(load)
            self.socket.send_pyobj(payload)
            data = self.auth.crypticle.loads(self.socket.recv_pyobj())
            if data == False:
                return False
            if not data:
                break
            fn_.write(data)
        return dest

    def cache_file(self, path, env='base'):
        '''
        Pull a file down from the file server and store it in the minion file
        cache
        '''
        dest = os.path.join(self.opts['cachedir'], 'files', path)
        return self.get_file(path, dest, True, env)

    def cache_files(self, paths, env='base'):
        '''
        Download a list of files stored on the master and put them in the minion
        file cache
        '''
        ret = []
        for path in paths:
            ret.append(self.cache_file(path, env))
        return ret

    def hash_file(self, path, env='base'):
        '''
        Return the hash of a file, to get the hash of a file on the
        salt master file server prepend the path with salt://<file on server>
        otherwise, prepend the file with / for a local file.
        '''
        path = self._check_proto(path)
        payload = {'enc': 'aes'}
        load = {'path': path,
                'env': env,
                'cmd': '_file_hash'}
        payload['load'] = self.auth.crypticle.dumps(load)
        self.socket.send_pyobj(payload)
        return self.auth.crypticle.loads(socket.recv_pyobj())

    def get_state(self, sls, env):
        '''
        Get a state file from the master and store it in the local minion cache
        return the location of the file
        '''
        pacman.cache = pacman/cache.sls or pacman/cache/init.sls
        pacman = pacman/init.sls
        if sls.count('.'):
            root = sls.replace('.', '/')
            for path in [root + '.sls', os.path.join(root, 'init.sls')]:
                dest = self.cache_file(path, env)
                if dest:
                    return dest
        return False

    def master_opts(self):
        '''
        Return the master opts data
        '''
        payload = {'enc': 'aes'}
        load = {'cmd': '_master_opts'}
        payload['load'] = self.auth.crypticle.dumps(load)
        self.socket.send_pyobj(payload)
        return self.auth.crypticle.loads(socket.recv_pyobj())

<commit_msg>Add method tot he mater for matching top file components<commit_after>'''
Routines to set up a minion
'''
# Import python libs
import os
import distutils.sysconfig
import glob
import re
import time
import logging
import tempfile
import traceback
import shutil
import threading
import multiprocessing

# Import zeromq libs
import zmq
# Import salt libs
import salt.crypt
from salt.crypt import AuthenticationError
import salt.utils
import salt.modules
import salt.returners
import salt.loader

log = logging.getLogger(__name__)

# To set up a minion:
# 1, Read in the configuration
# 2. Generate the function mapping dict
# 3. Authenticate with the master
# 4. Store the aes key
# 5. connect to the publisher
# 6. handle publications

class MinionError(Exception): pass

class Minion(object):
    '''
    This class instantiates a minion, runs connections for a minion, and loads
    all of the functions into the minion
    '''
    def __init__(self, opts):
        '''
        Pass in the options dict
        '''
        self.opts = opts
        self.mod_opts = self.__prep_mod_opts()
        self.functions, self.returners = self.__load_modules()
        self.matcher = Matcher(self.opts, self.functions)
        self.authenticate()

    def __prep_mod_opts(self):
        '''
        Returns a deep copy of the opts with key bits stripped out
        '''
        mod_opts = {}
        for key, val in self.opts.items():
            if key == 'logger':
                continue
            mod_opts[key] = val
        return mod_opts

    def __load_modules(self):
        '''
        Return the functions and the returners loaded up from the loader module
        '''
        functions = salt.loader.minion_mods(self.opts)
        returners = salt.loader.returners(self.opts)
        return functions, returners

    def _handle_payload(self, payload):
        '''
        Takes a payload from the master publisher and does whatever the
        master wants done.
        '''
        {'aes': self._handle_aes,
         'pub': self._handle_pub,
         'clear': self._handle_clear}[payload['enc']](payload['load'])

    def _handle_aes(self, load):
        '''
        Takes the aes encrypted load, decrypts is and runs the encapsulated
        instructions
        '''
        data = None
        try:
            data = self.crypticle.loads(load)
        except AuthenticationError:
            self.authenticate()
            data = self.crypticle.loads(load)
        # Verify that the publication is valid
        if not data.has_key('tgt')\
                or not data.has_key('jid')\
                or not data.has_key('fun')\
                or not data.has_key('arg'):
            return
        # Verify that the publication applies to this minion
        if data.has_key('tgt_type'):
            if not getattr(self.matcher, data['tgt_type'] + '_match')(data['tgt']):
                return
        else:
            if not self.matcher.glob_match(data['tgt']):
                return
        self._handle_decoded_payload(data)

    def _handle_pub(self, load):
        '''
        Handle public key payloads
        '''
        pass

    def _handle_clear(self, load):
        '''
        Handle un-encrypted transmissions
        '''
        pass

    def _handle_decoded_payload(self, data):
        '''
        Override this method if you wish to handle the decoded data differently.
        '''
        if self.opts['multiprocessing']:
            if type(data['fun']) == type(list()):
                multiprocessing.Process(
                    target=lambda: self._thread_multi_return(data)
                ).start()
            else:
                multiprocessing.Process(
                    target=lambda: self._thread_return(data)
                ).start()
        else:
            if type(data['fun']) == type(list()):
                threading.Thread(
                    target=lambda: self._thread_multi_return(data)
                ).start()
            else:
                threading.Thread(
                    target=lambda: self._thread_return(data)
                ).start()

    def _thread_return(self, data):
        '''
        This method should be used as a threading target, start the actual
        minion side execution.
        '''
        ret = {}
        for ind in range(0, len(data['arg'])):
            try:
                data['arg'][ind] = eval(data['arg'][ind])
            except:
                pass

        try:
            ret['return'] = self.functions[data['fun']](*data['arg'])
        except Exception as exc:
            trb = traceback.format_exc()
            log.warning('The minion function caused an exception: %s', exc)
            ret['return'] = trb
        ret['jid'] = data['jid']
        if data['ret']:
            ret['id'] = self.opts['id']
            try:
                self.returners[data['ret']](ret)
            except Exception as exc:
                log.error('The return failed for job %s %s', data['jid'], exc)
        else:
            self._return_pub(ret)

    def _thread_multi_return(self, data):
        '''
        This method should be used as a threading target, start the actual
        minion side execution.
        '''
        ret = {'return': {}}
        for ind in range(0, len(data['fun'])):
            for index in range(0, len(data['arg'][ind])):
                try:
                    data['arg'][ind][index] = eval(data['arg'][ind][index])
                except:
                    pass

            try:
                ret['return'][data['fun'][ind]]\
                    = self.functions[data['fun'][ind]](*data['arg'][ind])
            except Exception as exc:
                trb = traceback.format_exc()
                log.warning('The minion function caused an exception: %s', exc)
                ret['return'][data['fun'][ind]] = trb
            ret['jid'] = data['jid']
        if data['ret']:
            ret['id'] = self.opts['id']
            try:
                self.returners[data['ret']](ret)
            except Exception as exc:
                log.error('The return failed for job %s %s', data['jid'], exc)
        else:
            self._return_pub(ret)

    def _return_pub(self, ret):
        '''
        Return the data from the executed command to the master server
        '''
        log.info('Returning information for job: %(jid)s', ret)
        context = zmq.Context()
        socket = context.socket(zmq.REQ)
        socket.connect(self.opts['master_uri'])
        payload = {'enc': 'aes'}
        load = {'return': ret['return'],
                'cmd': '_return',
                'jid': ret['jid'],
                'id': self.opts['id']}
        payload['load'] = self.crypticle.dumps(load)
        socket.send_pyobj(payload)
        return socket.recv()

    def reload_functions(self):
        '''
        Reload the functions dict for this minion, reading in any new functions
        '''
        self.functions = self.__load_functions()
        log.debug('Refreshed functions, loaded functions: %s', self.functions)
        return True

    def authenticate(self):
        '''
        Authenticate with the master, this method breaks the functional
        paradigm, it will update the master information from a fresh sign in,
        signing in can occur as often as needed to keep up with the revolving
        master aes key.
        '''
        log.debug('Attempting to authenticate with the Salt Master')
        auth = salt.crypt.Auth(self.opts)
        while True:
            creds = auth.sign_in()
            if creds != 'retry':
                log.info('Authentication with master successful!')
                break
            log.info('Waiting for minion key to be accepted by the master.')
            time.sleep(10)
        self.aes = creds['aes']
        self.publish_port = creds['publish_port']
        self.crypticle = salt.crypt.Crypticle(self.aes)

    def tune_in(self):
        '''
        Lock onto the publisher. This is the main event loop for the minion
        '''
        master_pub = 'tcp://' + self.opts['master'] + ':'\
                   + str(self.publish_port)
        context = zmq.Context()
        socket = context.socket(zmq.SUB)
        socket.connect(master_pub)
        socket.setsockopt(zmq.SUBSCRIBE, '')
        while True:
            payload = socket.recv_pyobj()
            self._handle_payload(payload)

class Matcher(object):
    '''
    Use to return the value for matching calls from the master
    '''
    def __init__(self, opts, functions=None):
        self.opts = opts
        if not functions:
            functions = salt.loader.minion_mods(self.opts)
        else:
            self.functions = functions

    def confirm_top(self, data):
        '''
        Takes the data passed to a top file environment and determines if the
        data matches this minion
        '''
        matcher = 'glob'
        for item in data:
            if type(item) == type(dict()):
                if item.has_key('match'):
                    matcher = item['match']
        if hasattr(self, matcher + '_match'):
            return getattr(self, matcher)
        else:
            log.error('Attempting to match with unknown matcher: %s', matcher)
            return False

    def glob_match(self, tgt):
        '''
        Returns true if the passed glob matches the id
        '''
        tmp_dir = tempfile.mkdtemp()
        cwd = os.getcwd()
        os.chdir(tmp_dir)
        open(self.opts['id'], 'w+').write('salt')
        ret = bool(glob.glob(tgt))
        os.chdir(cwd)
        shutil.rmtree(tmp_dir)
        return ret

    def pcre_match(self, tgt):
        '''
        Returns true if the passed pcre regex matches
        '''
        return bool(re.match(tgt, self.opts['id']))

    def list_match(self, tgt):
        '''
        Determines if this host is on the list
        '''
        return bool(tgt.count(self.opts['id']))

    def grain_match(self, tgt):
        '''
        Reads in the grains regular expression match
        '''
        comps = tgt.split(':')
        return bool(re.match(comps[1], self.opts['grains'][comps[0]]))

    def exsel_match(self, tgt):
        '''
        Runs a function and return the exit code
        '''
        if not self.functions.has_key(tgt):
            return False
        return(self.functions[tgt]())

class FileClient(object):
    '''
    Interact with the salt master file server.
    '''
    def __init__(self, opts):
        self.opts = opts
        self.auth = salt.crypt.SAuth(opts)
        self.socket = self.__get_socket()

    def __get_socket(self):
        '''
        Return the ZeroMQ socket to use
        '''
        context = zmq.Context()
        socket = context.socket(zmq.REQ)
        socket.connect(self.opts['master_uri'])
        return socket

    def _check_proto(self, path):
        '''
        Make sure that this path is intended for the salt master and trim it
        '''
        if not path.startswith('salt://'):
            raise MinionError('Unsupported path')
        return path[:7]

    def get_file(self, path, dest, makedirs=False, env='base'):
        '''
        Get a single file from the salt-master
        '''
        path = self._check_proto(path)
        payload = {'enc': 'aes'}
        destdir = os.path.dirname(dest)
        if not os.path.isdir(destdir):
            if makedirs:
                os.makedirs(destdir)
            else:
                return False
        fn_ = open(dest, 'w+')
        load = {'path': path,
                'env': env,
                'cmd': '_serve_file'}
        while True:
            load['loc'] = fn_.tell()
            payload['load'] = self.crypticle.dumps(load)
            self.socket.send_pyobj(payload)
            data = self.auth.crypticle.loads(self.socket.recv_pyobj())
            if data == False:
                return False
            if not data:
                break
            fn_.write(data)
        return dest

    def cache_file(self, path, env='base'):
        '''
        Pull a file down from the file server and store it in the minion file
        cache
        '''
        dest = os.path.join(self.opts['cachedir'], 'files', path)
        return self.get_file(path, dest, True, env)

    def cache_files(self, paths, env='base'):
        '''
        Download a list of files stored on the master and put them in the minion
        file cache
        '''
        ret = []
        for path in paths:
            ret.append(self.cache_file(path, env))
        return ret

    def hash_file(self, path, env='base'):
        '''
        Return the hash of a file, to get the hash of a file on the
        salt master file server prepend the path with salt://<file on server>
        otherwise, prepend the file with / for a local file.
        '''
        path = self._check_proto(path)
        payload = {'enc': 'aes'}
        load = {'path': path,
                'env': env,
                'cmd': '_file_hash'}
        payload['load'] = self.auth.crypticle.dumps(load)
        self.socket.send_pyobj(payload)
        return self.auth.crypticle.loads(socket.recv_pyobj())

    def get_state(self, sls, env):
        '''
        Get a state file from the master and store it in the local minion cache
        return the location of the file
        '''
        pacman.cache = pacman/cache.sls or pacman/cache/init.sls
        pacman = pacman/init.sls
        if sls.count('.'):
            root = sls.replace('.', '/')
            for path in [root + '.sls', os.path.join(root, 'init.sls')]:
                dest = self.cache_file(path, env)
                if dest:
                    return dest
        return False

    def master_opts(self):
        '''
        Return the master opts data
        '''
        payload = {'enc': 'aes'}
        load = {'cmd': '_master_opts'}
        payload['load'] = self.auth.crypticle.dumps(load)
        self.socket.send_pyobj(payload)
        return self.auth.crypticle.loads(socket.recv_pyobj())

<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Skewed Student distribution.

.. [1] Hansen, B. E. (1994). Autoregressive conditional density estimation.
    International Economic Review, 35(3), 705–730.

"""

from __future__ import print_function, division

import numpy as np
import matplotlib.pylab as plt
import seaborn as sns

from scipy.special import gamma

__author__ = "Stanislav Khrapov"
__email__ = "khrapovs@gmail.com"

class SkewStudent(object):

    """Skewed Student distribution class.

    """

    def __init__(self, nup=10, lam=-.1):
        """Initialize the class.

        Parameters
        ----------
        nup : float
            Degrees of freedom
        lam : float
            Skewness

        """
        self.nup = nup
        self.lam = lam

    def pdf(self, arg):
        """Probability density function.

        """
        c = gamma((self.nup+1)/2)/((np.pi*(self.nup-2))**.5*gamma(self.nup/2))
        a = 4*self.lam*c*(self.nup-2)/(self.nup-1)
        b = (1 + 3*self.lam**2 - a**2)**.5

        pdf1 = b*c*(1 + 1/(self.nup-2)*((b*arg+a)/(1-self.lam))**2) \
            **(-(self.nup+1)/2)
        pdf2 = b*c*(1 + 1/(self.nup-2)*((b*arg+a)/(1+self.lam))**2) \
            **(-(self.nup+1)/2)

        return pdf1 * (arg < -a/b) + pdf2 * (arg >= -a/b)

    def plot_density(self, arg=np.linspace(-2, 2, 100)):
        """Plot probability density function.

        """
        plt.plot(arg, self.pdf(arg))
        plt.show()

if __name__ == '__main__':

    sns.set_context('notebook')
    skewt = SkewStudent()
    skewt.plot_density()
<commit_msg>Docstrings<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Skewed Student distribution.

.. [1] Hansen, B. E. (1994). Autoregressive conditional density estimation.
    International Economic Review, 35(3), 705–730.

"""

from __future__ import print_function, division

import numpy as np
import matplotlib.pylab as plt
import seaborn as sns

from scipy.special import gamma

__author__ = "Stanislav Khrapov"
__email__ = "khrapovs@gmail.com"

class SkewStudent(object):

    """Skewed Student distribution class.

    """

    def __init__(self, nup=10, lam=-.1):
        """Initialize the class.

        Parameters
        ----------
        nup : float
            Degrees of freedom
        lam : float
            Skewness

        """
        self.nup = nup
        self.lam = lam

    def pdf(self, arg):
        """Probability density function.

        Parameters
        ----------
        arg : array
            Grid of point to evaluate PDF at

        Returns
        -------
        pdf : array
            PDF values. Same shape as the input.

        """
        c = gamma((self.nup+1)/2)/((np.pi*(self.nup-2))**.5*gamma(self.nup/2))
        a = 4*self.lam*c*(self.nup-2)/(self.nup-1)
        b = (1 + 3*self.lam**2 - a**2)**.5

        pdf1 = b*c*(1 + 1/(self.nup-2)*((b*arg+a)/(1-self.lam))**2) \
            **(-(self.nup+1)/2)
        pdf2 = b*c*(1 + 1/(self.nup-2)*((b*arg+a)/(1+self.lam))**2) \
            **(-(self.nup+1)/2)

        return pdf1 * (arg < -a/b) + pdf2 * (arg >= -a/b)

    def plot_density(self, arg=np.linspace(-2, 2, 100)):
        """Plot probability density function.

        """
        plt.plot(arg, self.pdf(arg))
        plt.show()

if __name__ == '__main__':

    sns.set_context('notebook')
    skewt = SkewStudent()
    skewt.plot_density()
<|endoftext|>