diff --git a/docker/push.sh b/docker/push.sh
new file mode 100755
index 0000000..589d2c0
--- /dev/null
+++ b/docker/push.sh
@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+docker login
+
+docker build -t perara/wg-manager .
+docker push perara/wg-manager
diff --git a/wg_dashboard_backend/database.py b/wg_dashboard_backend/database.py
index b22a012..e55f707 100644
--- a/wg_dashboard_backend/database.py
+++ b/wg_dashboard_backend/database.py
@@ -11,3 +11,12 @@ engine = sqlalchemy.create_engine(
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 
 Base = declarative_base()
+
+from const import DATABASE_URL
+from migrate import DatabaseAlreadyControlledError
+from migrate.versioning.shell import main
+try:
+    main(["version_control", DATABASE_URL, "migrations"])
+except DatabaseAlreadyControlledError:
+    pass
+main(["upgrade", DATABASE_URL, "migrations"])
diff --git a/wg_dashboard_backend/db/wireguard.py b/wg_dashboard_backend/db/wireguard.py
index 4be3eef..4ee50c2 100644
--- a/wg_dashboard_backend/db/wireguard.py
+++ b/wg_dashboard_backend/db/wireguard.py
@@ -22,6 +22,10 @@ def start_client(sess: Session, peer: schemas.WGPeer):
     output = subprocess.check_output(const.CMD_WG_QUICK + ["up", client_file], stderr=subprocess.STDOUT)
 
 
+def get_server_by_id(sess: Session, server_id):
+    return sess.query(models.WGServer).filter_by(id=server_id).one()
+
+
 def peer_query_get_by_address(sess: Session, address: str, server: str):
     return sess.query(models.WGPeer) \
         .filter(models.WGPeer.address == address) \
diff --git a/wg_dashboard_backend/migrations/__init__.py b/wg_dashboard_backend/migrations/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/wg_dashboard_backend/migrations/versions/001_add_dns_column.py b/wg_dashboard_backend/migrations/versions/001_add_dns_column.py
new file mode 100644
index 0000000..bc7036a
--- /dev/null
+++ b/wg_dashboard_backend/migrations/versions/001_add_dns_column.py
@@ -0,0 +1,17 @@
+from sqlalchemy import Table, MetaData, String, Column, Text
+
+
+def upgrade(migrate_engine):
+    try:
+        meta = MetaData(bind=migrate_engine)
+        account = Table('peer', meta, autoload=True)
+        dns = Column('dns', Text)
+        dns.create(account)
+    except:
+        pass
+
+
+def downgrade(migrate_engine):
+    meta = MetaData(bind=migrate_engine)
+    dns = Table('peer', meta, autoload=True)
+    dns.c.email.drop()
diff --git a/wg_dashboard_backend/migrations/versions/002_remove_server_shared_key.py b/wg_dashboard_backend/migrations/versions/002_remove_server_shared_key.py
new file mode 100644
index 0000000..c6215a3
--- /dev/null
+++ b/wg_dashboard_backend/migrations/versions/002_remove_server_shared_key.py
@@ -0,0 +1,16 @@
+from sqlalchemy import *
+from migrate import *
+
+
+def upgrade(migrate_engine):
+    meta = MetaData(bind=migrate_engine)
+    server = Table('server', meta, autoload=True)
+    server.c.shared_key.drop()
+
+
+def downgrade(migrate_engine):
+    meta = MetaData(bind=migrate_engine)
+    account = Table('server', meta, autoload=True)
+    shared_key = Column('shared_key', Text)
+    shared_key.create(account)
+
diff --git a/wg_dashboard_backend/migrations/versions/003_create_client_shared_key.py b/wg_dashboard_backend/migrations/versions/003_create_client_shared_key.py
new file mode 100644
index 0000000..12ad910
--- /dev/null
+++ b/wg_dashboard_backend/migrations/versions/003_create_client_shared_key.py
@@ -0,0 +1,19 @@
+from sqlalchemy import *
+from migrate import *
+
+
+def upgrade(migrate_engine):
+    try:
+        meta = MetaData(bind=migrate_engine)
+        peer = Table('peer', meta, autoload=True)
+        shared_key = Column('shared_key', Text)
+        shared_key.create(peer)
+    except:
+        pass
+
+
+def downgrade(migrate_engine):
+    meta = MetaData(bind=migrate_engine)
+    dns = Table('peer', meta, autoload=True)
+    dns.c.shared_key.drop()
+
diff --git a/wg_dashboard_backend/migrations/versions/__init__.py b/wg_dashboard_backend/migrations/versions/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/wg_dashboard_backend/models.py b/wg_dashboard_backend/models.py
index 4de51c4..0e4a07d 100644
--- a/wg_dashboard_backend/models.py
+++ b/wg_dashboard_backend/models.py
@@ -1,7 +1,7 @@
 import sqlalchemy
+
 from sqlalchemy import Integer, Column
 from sqlalchemy.orm import relationship, backref
-
 from database import Base
 
 
@@ -25,7 +25,6 @@ class WGServer(Base):
     listen_port = Column(sqlalchemy.String, unique=True)
     private_key = Column(sqlalchemy.String)
     public_key = Column(sqlalchemy.String)
-    shared_key = Column(sqlalchemy.String)
     endpoint = Column(sqlalchemy.String)
     dns = Column(sqlalchemy.String)
 
@@ -45,11 +44,10 @@ class WGPeer(Base):
     address = Column(sqlalchemy.String)
     public_key = Column(sqlalchemy.String)
     private_key = Column(sqlalchemy.String)
-    dns = Column(sqlalchemy.String)
+    shared_key = Column(sqlalchemy.Text)
+    dns = Column(sqlalchemy.Text)
     allowed_ips = Column(sqlalchemy.String)
 
     server_id = Column(Integer, sqlalchemy.ForeignKey('server.id', ondelete="CASCADE", onupdate="CASCADE"))
     server = relationship("WGServer", backref=backref("server"))
     configuration = Column(sqlalchemy.Text)
-
-
diff --git a/wg_dashboard_backend/requirements.txt b/wg_dashboard_backend/requirements.txt
index 299dd4a..1ec17e6 100644
--- a/wg_dashboard_backend/requirements.txt
+++ b/wg_dashboard_backend/requirements.txt
@@ -10,4 +10,5 @@ bcrypt
 python-multipart
 jinja2
 sqlalchemy_utils
+sqlalchemy-migrate
 requests
diff --git a/wg_dashboard_backend/routers/v1/peer.py b/wg_dashboard_backend/routers/v1/peer.py
index dbef6fc..c81aa08 100644
--- a/wg_dashboard_backend/routers/v1/peer.py
+++ b/wg_dashboard_backend/routers/v1/peer.py
@@ -87,22 +87,32 @@ def delete_peer(
     return peer
 
 
-@router.post("/edit", response_model=schemas.WGPeer)
+@router.post("/edit")
 def edit_peer(
         peer: schemas.WGPeer,
         sess: Session = Depends(middleware.get_db)
 ):
-    server = schemas.WGServer(interface="")\
-            .from_orm(sess.query(models.WGServer).filter_by(id=peer.server_id).one())
+    # Retrieve server from db
+    server: models.WGServer = db.wireguard.get_server_by_id(sess, peer.server_id)
 
+    # Generate peer configuration
     peer.configuration = script.wireguard.generate_config(dict(
         peer=peer,
         server=server
     ))
+
+    # Update database record for Peer
     sess.query(models.WGPeer)\
         .filter_by(id=peer.id)\
         .update(peer.dict(exclude={"id"}))
 
+    # Generate server configuration
+    server.configuration = script.wireguard.generate_config(server)
+    sess.add(server)
+
     sess.commit()
 
-    return peer
+    return dict(
+        peer=peer,
+        server_configuration=server.configuration
+    )
diff --git a/wg_dashboard_backend/schemas.py b/wg_dashboard_backend/schemas.py
index f37fa9e..6a53614 100644
--- a/wg_dashboard_backend/schemas.py
+++ b/wg_dashboard_backend/schemas.py
@@ -116,6 +116,7 @@ class WGPeer(GenericModel):
     address: str = None
     private_key: str = None
     public_key: str = None
+    shared_key: str = None
     server_id: str
     dns: str = None
     allowed_ips: str = None
@@ -148,7 +149,6 @@ class WGServer(GenericModel):
     endpoint: str = None
     private_key: str = None
     public_key: str = None
-    shared_key: str = None
     is_running: bool = None
     configuration: str = None
     post_up: str = None
diff --git a/wg_dashboard_backend/script/wireguard.py b/wg_dashboard_backend/script/wireguard.py
index 3f57687..53bbbaa 100644
--- a/wg_dashboard_backend/script/wireguard.py
+++ b/wg_dashboard_backend/script/wireguard.py
@@ -5,6 +5,7 @@ import tempfile
 import typing
 
 import const
+import models
 import schemas
 import os
 import re
@@ -195,7 +196,7 @@ def move_server_dir(interface, interface1):
 def generate_config(obj: typing.Union[typing.Dict[schemas.WGPeer, schemas.WGServer], schemas.WGServer]):
     if isinstance(obj, dict) and "server" in obj and "peer" in obj:
         template = "peer.j2"
-    elif isinstance(obj, schemas.WGServer):
+    elif isinstance(obj, schemas.WGServer) or isinstance(obj, models.WGServer):
         template = "server.j2"
     else:
         raise ValueError("Incorrect input type. Should be WGPeer or WGServer")
diff --git a/wg_dashboard_backend/templates/peer.j2 b/wg_dashboard_backend/templates/peer.j2
index fdd19a5..c02656a 100644
--- a/wg_dashboard_backend/templates/peer.j2
+++ b/wg_dashboard_backend/templates/peer.j2
@@ -7,6 +7,6 @@ DNS = {{ data.peer.dns }}
 PublicKey = {{ data.server.public_key }}
 AllowedIPs = {{ data.peer.allowed_ips }}
 Endpoint = {{ data.server.endpoint }}:{{ data.server.listen_port }}
-{% if data.server.shared_key %}
-PresharedKey = {{ data.server.shared_key }}
+{% if data.peer.shared_key %}
+PresharedKey = {{ data.peer.shared_key }}
 {% endif %}
diff --git a/wg_dashboard_backend/templates/server.j2 b/wg_dashboard_backend/templates/server.j2
index 123f94e..e4f50dd 100644
--- a/wg_dashboard_backend/templates/server.j2
+++ b/wg_dashboard_backend/templates/server.j2
@@ -10,8 +10,8 @@ PostDown = {{ data.post_down }}
 [Peer]
 # Client Name: {{ peer.name }}
 PublicKey = {{ peer.public_key }}
-{%- if peer.preshared_key -%}
-PresharedKey = {{ peer.preshared_key }}
-{% endif %}
+{%- if peer.shared_key %}
+PresharedKey = {{ peer.shared_key }}
+{%- endif %}
 AllowedIPs = {{ peer.address }}
 {% endfor %}
diff --git a/wg_dashboard_frontend/src/app/interfaces/peer.ts b/wg_dashboard_frontend/src/app/interfaces/peer.ts
index 99d61c0..f8ca8aa 100644
--- a/wg_dashboard_frontend/src/app/interfaces/peer.ts
+++ b/wg_dashboard_frontend/src/app/interfaces/peer.ts
@@ -3,6 +3,7 @@ export interface Peer {
   address: string;
   public_key: string;
   private_key: string;
+  shared_key: string;
   dns: string;
   allowed_ips: string;
   name: string;
diff --git a/wg_dashboard_frontend/src/app/page/dashboard/add-server/add-server.component.html b/wg_dashboard_frontend/src/app/page/dashboard/add-server/add-server.component.html
index 0f6202c..a30a80a 100644
--- a/wg_dashboard_frontend/src/app/page/dashboard/add-server/add-server.component.html
+++ b/wg_dashboard_frontend/src/app/page/dashboard/add-server/add-server.component.html
@@ -74,22 +74,12 @@
         </mat-form-field>
       </p>
 
-      <p>
-        <mat-form-field class="add-server-full-width">
-          <mat-label>Shared-Key</mat-label>
-          <input formControlName="shared_key" matInput>
-        </mat-form-field>
-      </p>
-
       <div class="button-row">
         <button type="button" [disabled]="!isEdit" (click)="getKeyPair()" mat-raised-button color="primary">
           <i class="material-icons">vpn_key</i>
           Generate KeyPair
         </button>
-        <button type="button" [disabled]="!isEdit" (click)="getPSK()" mat-raised-button color="primary">
-          <i class="material-icons">share</i>
-          Generate PSK
-        </button>
+
       </div>
 
       <p><b>Scripts</b></p>
diff --git a/wg_dashboard_frontend/src/app/page/dashboard/add-server/add-server.component.ts b/wg_dashboard_frontend/src/app/page/dashboard/add-server/add-server.component.ts
index 84530c8..ab81121 100644
--- a/wg_dashboard_frontend/src/app/page/dashboard/add-server/add-server.component.ts
+++ b/wg_dashboard_frontend/src/app/page/dashboard/add-server/add-server.component.ts
@@ -47,7 +47,6 @@ export class AddServerComponent implements OnInit {
       dns: new FormControl(''),
       private_key: new FormControl('' ),
       public_key: new FormControl('' ),
-      shared_key: new FormControl('' ),
       post_up: new FormControl(''),
       post_down: new FormControl(''),
 
@@ -216,13 +215,7 @@ export class AddServerComponent implements OnInit {
     });
   }
 
-  getPSK() {
-    this.serverAPI.getPSK().subscribe((psk: any) => {
-      this.serverForm.patchValue({
-        shared_key: psk.psk,
-      });
-    });
-  }
+
 
   resetForm() {
     this.isEdit = false;
diff --git a/wg_dashboard_frontend/src/app/page/dashboard/peer/peer.component.html b/wg_dashboard_frontend/src/app/page/dashboard/peer/peer.component.html
index 564fd95..bfca94a 100644
--- a/wg_dashboard_frontend/src/app/page/dashboard/peer/peer.component.html
+++ b/wg_dashboard_frontend/src/app/page/dashboard/peer/peer.component.html
@@ -56,6 +56,29 @@
         </mat-form-field>
       </p>
 
+
+      <table class="full-width" cellspacing="0"><tr>
+        <td width="80%">
+          <mat-form-field class="full-width">
+            <mat-label>PreShared-Key</mat-label>
+            <input [disabled]="!peer._edit" name="shared_key" [(ngModel)]="peer.shared_key" matInput>
+          </mat-form-field>
+        </td>
+        <td>
+          <button type="button" [disabled]="!peer._edit" (click)="getPSK()" mat-raised-button color="primary">
+            <i class="material-icons">share</i>
+            Generate PSK
+          </button>
+        </td>
+      </tr></table>
+
+      <p>
+
+      </p>
+
+
+
+
       <button
         [hidden]="!peer._edit"
         [disabled]="!peerForm.valid"
diff --git a/wg_dashboard_frontend/src/app/page/dashboard/peer/peer.component.ts b/wg_dashboard_frontend/src/app/page/dashboard/peer/peer.component.ts
index 7d8d3d4..fc50bf3 100644
--- a/wg_dashboard_frontend/src/app/page/dashboard/peer/peer.component.ts
+++ b/wg_dashboard_frontend/src/app/page/dashboard/peer/peer.component.ts
@@ -41,9 +41,10 @@ export class PeerComponent implements OnInit {
 
       // Submit the edit (True -> False)
       const idx = this.server.peers.indexOf(this.peer);
-      this.serverAPI.editPeer(this.peer).subscribe((newPeer) => {
-        Object.keys(newPeer).forEach(k => {
-          this.server.peers[idx][k] = newPeer[k];
+      this.serverAPI.editPeer(this.peer).subscribe((data) => {
+        this.server.configuration = data.server_configuration;
+        Object.keys(data.peer).forEach(k => {
+          this.server.peers[idx][k] = data.peer[k];
         });
       });
 
@@ -65,7 +66,10 @@ export class PeerComponent implements OnInit {
     });
   }
 
-
-
+  getPSK() {
+    this.serverAPI.getPSK().subscribe((psk: any) => {
+      this.peer.shared_key = psk.psk;
+    });
+  }
 
 }
diff --git a/wg_dashboard_frontend/src/app/services/server.service.ts b/wg_dashboard_frontend/src/app/services/server.service.ts
index db8fa00..2e61b2c 100644
--- a/wg_dashboard_frontend/src/app/services/server.service.ts
+++ b/wg_dashboard_frontend/src/app/services/server.service.ts
@@ -35,7 +35,7 @@ export class ServerService {
     return this.http.post(this.peerURL + '/add', server_interface);
   }
 
-  public editPeer(peer: Peer): Subscribable<Peer> {
+  public editPeer(peer: Peer): Subscribable<{ peer: Peer, server_configuration: string }> {
     return this.http.post(this.peerURL + '/edit', peer);
   }