^^&&^^

Edit File: r1soft.py

import json
import os
import re
import socket
import subprocess
import time
from json import JSONDecodeError

import requests

from restore_infected.backup_backends_lib import (
    BackupBase,
    backend_auth_required,
    extra,
)
from restore_infected.helpers import DateTime, from_env

TOKEN_FILE = '/var/restore_infected/r1soft_api_token.json'

auth_required = backend_auth_required(TOKEN_FILE, "Initialize R1Soft first!")

def is_suitable():
    return True

class R1SoftConnector:
    api_port = 9080
    api_addr = 'http://{}:{}/rest'
    api_get_token = 'user/authenticate'

agent_cached = None
    machine_cached = None

class InternalServerError(Exception):
        def __init__(self, url):
            message = 'Internal server error. {}'.format(url)
            super().__init__(message)

class ConnectionError(Exception):
        def __init__(self, url, status_code, content):
            message = 'Request to {} failed ({}): {}' \
                .format(url, status_code, content)
            super().__init__(message)

def __init__(self, ip, encryption_key):
        self.ip = ip
        self.__encryption_key = encryption_key

def save_token(self, username, password):
        """
        Receives the auth token from the server and saves it in the file.
        :param username: R1Soft server username
        :param password: R1Soft server password
        """
        s = requests.Session()
        s.auth = (username, password)
        url = self._build_api_url(self.ip, self.api_get_token)
        r = s.get(url)
        self._check_response(r)

token = r.json()['authToken']

config = {
            'ip': self.ip,
            'token': token,
            'encryption_key': self.__encryption_key,
            'username': username,
            'timestamp': int(time.time())
        }

self.write_token(config)

def refresh_token(self):
        r = self._api_request(requests.get, self.api_get_token)
        token = r['authToken']
        timestamp = int(time.time())
        self._update_token_file(token=token, timestamp=timestamp)

@classmethod
    def read_token(cls):
        with open(TOKEN_FILE) as t_file:
            return json.load(t_file)

@classmethod
    def write_token(cls, config):
        if os.path.dirname(TOKEN_FILE):
            os.makedirs(os.path.dirname(TOKEN_FILE), exist_ok=True)
        with open(TOKEN_FILE, 'w') as t_file:
            json.dump(config, t_file)

def remove_token(self):
        try:
            os.remove(TOKEN_FILE)
        except FileNotFoundError:
            pass

@classmethod
    def from_token(cls):
        """
        Reads ip and encryption key from the file created by
          save_token().
        Raises an exception if the file does not exists.
        :return: new R1SoftConnector with token field not None
        """
        if not os.path.exists(TOKEN_FILE):
            raise Exception('No auth token found. Get token first.')

with open(TOKEN_FILE) as t_file:
            config = json.load(t_file)
            ip = config['ip']
            encryption_key = config['encryption_key']

_cls = cls(ip, encryption_key)

return _cls

@staticmethod
    def _check_response(response):
        """
        Exit with error status unless response code is 200
        :param response: obj -> response object
        """
        if response.status_code == 500:
            raise R1SoftConnector.InternalServerError(response.url)

if response.status_code < 200 or response.status_code >= 400:
            raise R1SoftConnector.ConnectionError(response.url,
                                                  response.status_code,
                                                  response.content)

@classmethod
    def _build_api_url(cls, ip, api_path):
        api_addr = cls.api_addr.format(ip, cls.api_port)

return '{}/{}'.format(api_addr, api_path)

def _update_token_file(self, **kwargs):
        config = self.read_token()
        for key, value in kwargs.items():
            config[key] = value

self.write_token(config)

@staticmethod
    def _get_machine_address():
        hostname = socket.gethostname()

ip_a = subprocess.check_output(['ip', '-o', '-4', 'address', 'show'])
        ip_list = ip_a.decode('utf-8').splitlines()
        ip_pattern = re.compile(r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}')
        ip_list = [ip_pattern.findall(ip)[0] for ip in ip_list]

return ip_list, hostname

def _api_request(self, method, api_path, data=None):
        """
        Sends an API request.
        :param method: e.g. requests.get
        :param api_path: The endpoint path that goes after rest/ (see api_addr)
        :param data: json parameters
        :return: response as dict if possible, as plain text otherwise
        """
        url = self._build_api_url(self.ip, api_path)
        token = self.read_token()['token']

headers = {
            'AuthToken': token
        }

r = method(url, headers=headers, json=data)
        self._check_response(r)

try:
            res = r.json()
        except JSONDecodeError:
            res = r.text

return res

def _get_agents(self):
        return self._api_request(requests.get, 'agent')

def _get_disk_safes(self):
        return self._api_request(requests.get, 'disksafe')

def _get_machines(self):
        return self._api_request(requests.get, 'machine')

def _get_inodes(self, recovery_point, path):
        machine = self.get_machine()
        data = {
            'passphrase': self.__encryption_key,
            'basePath': path
        }

api_path = 'backup/{}/{}/file/id' \
            .format(machine['id'], recovery_point['recoveryPointID'])

return self._api_request(requests.post, api_path, data)

def _get_restore_history(self):
        """
        Retrieves restore history.
        R1Soft server may fail to generate history while starting restore
        process.
        :return: list of restore attempts on success, otherwise empty list.
        """
        try:
            history = self._api_request(requests.get, 'restore/file')

return history
        except R1SoftConnector.InternalServerError:
            return []

def _restore_completed(self, recovery_id):
        for recovery_entry in self._get_restore_history():
            if recovery_entry['id'] == recovery_id:
                return recovery_entry['recoveryStatus'] == 'FINISHED'

return False

def _find_this_agent(self, agents):
        ip_list, hostname = self._get_machine_address()
        addr_list = ip_list + [hostname]

if agents:
            key = 'hostname' if 'hostname' in agents[0] else 'hostnameIp'
            for agent in agents:
                if agent[key] in addr_list:
                    return agent

raise Exception('Agent with any of the addresses {} not found.'
                        .format(str(addr_list)))

def get_agent(self):
        if self.agent_cached is None:
            agents = self._get_agents()
            self.agent_cached = self._find_this_agent(agents)

return self.agent_cached

def get_disk_safes(self):
        agent_id = self.get_agent()['id']
        disk_safes_all = self._get_disk_safes()
        disk_safes = []
        for disk_safe in disk_safes_all:
            if disk_safe['agentID'] == agent_id:
                disk_safes.append(disk_safe)

return disk_safes

def get_recovery_points(self, disk_safe):
        disk_safe_id = disk_safe['id']

return self._api_request(
            requests.get, 'recoverypoint/{}/usable'.format(disk_safe_id))

def restore_file(self, recovery_point, path, dst):
        just_path = os.path.dirname(path)
        just_name = os.path.basename(path)
        final_path = os.path.join(dst, just_path.strip('/'))
        final_file_name = os.path.join(final_path, just_name)
        machine_id = self.get_machine()['id']
        rec_id = recovery_point['recoveryPointID']
        data = {
            'basePath': just_path,
            'restoreMethod': 'ALTERNATE',
            'restoreToMachineId': machine_id,
            'restoreToPath': final_path,
            'childTokens': [just_name],
            'passphrase': self.__encryption_key
        }
        api_path = 'restore/file/{}/{}'.format(machine_id, rec_id)

restore_id = self._api_request(requests.post, api_path, data)
        while not self._restore_completed(restore_id):
            pass

return final_file_name

def get_machine(self):
        if self.machine_cached is None:
            machines = self._get_machines()
            self.machine_cached = self._find_this_agent(machines)

return self.machine_cached

def get_file_entry(self, recovery_point, file_path):
        just_path = os.path.dirname(file_path)
        just_name = os.path.basename(file_path)
        machine = self.get_machine()
        inodes = self._get_inodes(recovery_point, just_path)
        inode = inodes.get(just_name, None)
        data = {
            'passphrase': self.__encryption_key,
            'basePath': just_path,
            'inodeNumbers': [inode]
        }

if not inode:
            raise Exception(
                'No backup for \'{}\' found'.format(file_path))

api_path = 'backup/{}/{}/file' \
            .format(machine['id'], recovery_point['recoveryPointID'])

return self._api_request(requests.post, api_path, data)[just_name]

class R1SoftBackup(BackupBase):
    """
    R1Soft backup entry
    """

def __init__(self, rec_point):
        self.rec_point = rec_point
        self.r1 = R1SoftConnector.from_token()
        super().__init__('', DateTime.fromtimestamp(
            rec_point['createdOnTimestampInMillis'] / 1000))

def __repr__(self):
        return json.dumps(self.rec_point)

def __str__(self):
        return self.__repr__()

def close(self):
        pass

def file_data(self, path):
        file_entry = self.r1.get_file_entry(self.rec_point, path)
        return FileData(
            path,
            DateTime.fromtimestamp(file_entry["modifyTime"] / 1000),
            file_entry["fileSize"],
        )

def restore(self, items, destination_folder="/tmp"):
        return {
            self.r1.restore_file(
                self.rec_point, item.filename, destination_folder
            ): item.filename
            for item in items
        }

class FileData:
    """
    R1Soft FileData entry
    """

def __init__(self, path, mtime, size):
        self.filename = path
        self.mtime = mtime
        self.size = size

def __str__(self):
        return '{} [{} bytes] {}'.format(self.mtime, self.size, self.filename)

@from_env(
    ip="IP",
    username="ACCOUNT_NAME",
    password="PASSWORD",
    encryption_key="ENCRYPTION_KEY",
)
def init(
    ip, username, password, encryption_key,
):
    r1 = R1SoftConnector(ip, encryption_key)
    r1.save_token(username, password)

@auth_required
def backups(until=None, tmp_dir=None):
    r1 = R1SoftConnector.from_token()
    disks = r1.get_disk_safes()
    recs = r1.get_recovery_points(disks[0])
    backup_list = []
    for rec_point in recs:
        backup = R1SoftBackup(rec_point)
        if until is None or backup.created >= until:
            backup_list.append(backup)

return backup_list

@auth_required
def info():
    return R1SoftConnector.read_token()

@auth_required
@extra
def refresh_token():
    r1 = R1SoftConnector.from_token()
    r1.refresh_token()